iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Support --stream-interval and per-request stream_interval (#60589)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Bugen Zhao
2026-10-09 14:10:42 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent ffdde24f54
commit 35e920d671
25 changed files with 1040 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/cmd/src/cli.rs
+14
View File
@@ -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<i32>,
/// 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()])
}
rust/src/cmd/src/cli/tests.rs
+34
View File
@@ -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,
rust/src/cmd/src/cli/unsupported.rs
-4
View File
@@ -359,10 +359,6 @@ pub struct EngineUnsupportedArgs {
#[arg(long)]
pub collect_detailed_traces: Option<Unsupported>,
/// The interval (or buffer size) for streaming in terms of token length.
#[arg(long)]
pub stream_interval: Option<Unsupported>,
/// 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.
rust/src/engine-core-client/examples/external_engine_logprobs.rs
+2 -1
View File
@@ -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<EngineCoreStreamOutput> {
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() {
rust/src/engine-core-client/src/client.rs
+18 -3
View File
@@ -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<EngineCoreOutputStream> {
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
rust/src/engine-core-client/src/client/imp.rs
+40 -33
View File
@@ -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<String>,
data_parallel_rank: Option<u32>,
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<Item = &'a EngineCoreOutput>,
) -> Vec<Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>>> {
self.request_reg.lock().senders_for_outputs(outputs)
outputs: impl IntoIterator<Item = EngineCoreStreamOutput>,
) -> Vec<EngineCoreStreamOutput> {
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<Item = &'a String>,
) -> Vec<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
) -> Vec<OutputSender> {
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
rust/src/engine-core-client/src/client/output_channel.rs
+321
View File
@@ -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<Result<Delivery>>;
/// Sending half of a request's output channel, owned by its registry entry.
#[derive(Debug)]
pub(crate) struct OutputSender {
tx: mpsc::UnboundedSender<Result<Delivery>>,
/// 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<EngineCoreStreamOutput>,
/// 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<Vec<u32>> {
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<Vec<Vec<u32>>> {
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]);
}
}
rust/src/engine-core-client/src/client/state.rs
+218 -68
View File
@@ -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<Result<EngineCoreStreamOutput>>;
pub type OutputReceiver = mpsc::UnboundedReceiver<Result<EngineCoreStreamOutput>>;
pub type UtilitySender = oneshot::Sender<Result<UtilityOutput>>;
pub type UtilityReceiver = oneshot::Receiver<Result<UtilityOutput>>;
@@ -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<String>,
data_parallel_rank: Option<u32>,
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<OutputSender> {
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<EngineCoreStreamOutput> {
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<Item = &'a EngineCoreOutput>,
) -> Vec<Option<OutputSender>> {
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<Item = &'a String>,
@@ -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<String> {
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<String> {
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<u32>, Option<EngineCoreFinishReason>)> {
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 {
rust/src/engine-core-client/src/client/stream.rs
+79 -4
View File
@@ -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<EngineCoreStreamOutput>),
}
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::iter::Once<EngineCoreStreamOutput>, std::vec::IntoIter<EngineCoreStreamOutput>>;
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<Item = Result<EngineCoreStreamOutput>> + 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<EngineCoreStreamOutput>;
type Item = Result<EngineCoreStreamDelivery>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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!(
rust/src/engine-core-client/src/lib.rs
+1 -1
View File
@@ -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};
rust/src/engine-core-client/src/protocol/sampling.rs
+7
View File
@@ -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<NonZeroU32>,
}
impl EngineCoreSamplingParams {
@@ -195,6 +201,7 @@ impl EngineCoreSamplingParams {
skip_reading_prefix_cache: None,
extra_args: None,
routed_experts_prompt_start: 0,
stream_interval: None,
}
}
}
rust/src/engine-core-client/src/tests/client.rs
+22 -21
View File
@@ -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::<Vec<_>>().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::<Vec<_>>().await;
let outputs = client.call(request).await.unwrap().into_outputs().collect::<Vec<_>>().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(
rust/src/llm/src/lib.rs
+22 -1
View File
@@ -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<StatsLogger>,
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<GenerateOutputStream> {
pub async fn generate(&self, mut req: GenerateRequest) -> Result<GenerateOutputStream> {
// 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
rust/src/llm/src/output.rs
+66 -13
View File
@@ -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<GenerateOutput>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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<GenerateOutput> {
// 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<GenerateOutput> {
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<GenerateOutput>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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)))
}
}
rust/src/llm/tests/generate.rs
+106
View File
@@ -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<EngineCoreFinishReason>); 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();
rust/src/mock-engine/src/tests.rs
+7 -3
View File
@@ -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);
rust/src/server/examples/external_engine_openai_qwen.rs
+2
View File
@@ -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,
rust/src/server/src/config.rs
+5
View File
@@ -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<i32>,
/// 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.
rust/src/server/src/lib.rs
+3 -1
View File
@@ -123,7 +123,9 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
.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)
rust/src/server/src/routes/openai/chat_completions/convert.rs
+13
View File
@@ -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<ChatToolChoice, ApiEr
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::num::NonZeroU32;
use std::sync::Arc;
use axum::http::{HeaderMap, StatusCode};
@@ -540,6 +542,15 @@ mod tests {
}
}
#[test]
fn chat_http_request_rejects_zero_stream_interval() {
let request = json!({
"messages": [{"role": "user", "content": "hello"}],
"stream_interval": 0,
});
assert!(serde_json::from_value::<ChatCompletionRequest>(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);
rust/src/server/src/routes/openai/chat_completions/types.rs
+8
View File
@@ -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<RepetitionDetectionParams>,
/// 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<NonZeroU32>,
}
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,
}
}
}
rust/src/server/src/routes/openai/completions/convert.rs
+9 -1
View File
@@ -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]
rust/src/server/src/routes/openai/completions/types.rs
+7
View File
@@ -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<RepetitionDetectionParams>,
/// 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<NonZeroU32>,
/// Length penalty for beam search
pub length_penalty: Option<f32>,
rust/src/text/src/lower.rs
+29
View File
@@ -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(&params)?;
validate_vocab_range(&params, &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(&params);
@@ -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(&params);
@@ -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(&params);
@@ -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(&params);
@@ -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(&params);
@@ -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(&params);
@@ -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 {
rust/src/text/src/request.rs
+7
View File
@@ -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<bool>,
/// Additional request parameters for custom extensions.
pub vllm_xargs: Option<HashMap<String, Value>>,
/// 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<NonZeroU32>,
}
#[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,
}
}
}