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

[Rust Frontend] Populate cached_token_count in responses (#44887)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-06-10 20:50:05 -07:00
committed by GitHub
parent 85a0ffae42
commit 5d5591d99b
34 changed files with 556 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/chat/examples/external_engine_chat_qwen.rs
+2 -2
View File
@@ -131,13 +131,13 @@ async fn main() -> Result<()> {
ChatEvent::LogprobsDelta { .. } => {}
ChatEvent::Done {
message,
output_token_count,
usage,
finish_reason: reason,
..
} => {
final_reasoning = message.reasoning().unwrap_or_default();
final_text = message.text();
final_output_token_count = output_token_count;
final_output_token_count = usage.output_token_count;
finish_reason = Some(reason);
break;
}
rust/src/chat/src/event.rs
+2 -5
View File
@@ -2,6 +2,7 @@ use std::ops::Deref;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use vllm_llm::TokenUsage;
use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
use crate::FinishReason;
@@ -197,11 +198,7 @@ pub enum ChatEvent {
/// metadata.
Done {
message: AssistantMessage,
/// Number of prompt tokens actually sent to the engine after chat
/// template rendering and tokenization.
prompt_token_count: usize,
/// Number of output tokens generated.
output_token_count: usize,
usage: TokenUsage,
finish_reason: FinishReason,
/// Connector-specific KV transfer parameters for disaggregated serving.
kv_transfer_params: Option<serde_json::Value>,
rust/src/chat/src/output/default/reasoning.rs
+11 -6
View File
@@ -178,8 +178,7 @@ pub(crate) async fn reasoning_event_stream(
y.yield_ok(next).await;
}
y.yield_ok(ContentEvent::Done {
prompt_token_count: finished.prompt_token_count,
output_token_count: finished.output_token_count,
usage: finished.usage,
finish_reason: finished.finish_reason,
kv_transfer_params: finished.kv_transfer_params,
})
@@ -289,8 +288,11 @@ mod tests {
token_ids: vec![],
logprobs: None,
finished: Some(vllm_text::Finished {
prompt_token_count: 3,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 3,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -322,8 +324,11 @@ mod tests {
delta: "def".to_string(),
},
ContentEvent::Done {
prompt_token_count: 3,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 3,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
},
rust/src/chat/src/output/default/tool.rs
+42 -20
View File
@@ -240,8 +240,7 @@ pub(crate) async fn tool_event_stream(
.await;
}
ContentEvent::Done {
prompt_token_count,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
} => {
@@ -250,8 +249,7 @@ pub(crate) async fn tool_event_stream(
}
y.yield_ok(AssistantEvent::Done {
prompt_token_count,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
})
@@ -465,8 +463,11 @@ mod tests {
})
})
.chain(std::iter::once(Ok(ContentEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
})));
@@ -506,8 +507,11 @@ mod tests {
delta: "ignored".to_string(),
}),
Ok(ContentEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -659,8 +663,11 @@ mod tests {
delta: "def".to_string(),
}),
Ok(ContentEvent::Done {
prompt_token_count: 3,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 3,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -697,8 +704,11 @@ mod tests {
delta: "def".to_string(),
},
AssistantEvent::Done {
prompt_token_count: 3,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 3,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
},
@@ -739,8 +749,11 @@ mod tests {
token_ids: vec![],
}),
Ok(ContentEvent::Done {
prompt_token_count: 1,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -779,8 +792,11 @@ mod tests {
token_ids: vec![],
},
AssistantEvent::Done {
prompt_token_count: 1,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
},
@@ -796,8 +812,11 @@ mod tests {
delta: "ignored".to_string(),
}),
Ok(ContentEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -901,8 +920,11 @@ mod tests {
delta: "ignored".to_string(),
}),
Ok(ContentEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
rust/src/chat/src/output/harmony/mod.rs
+1 -2
View File
@@ -366,8 +366,7 @@ async fn harmony_assistant_event_stream(
if let Some(finished) = finished {
y.yield_ok(AssistantEvent::Done {
prompt_token_count: finished.prompt_token_count,
output_token_count: finished.output_token_count,
usage: finished.usage,
finish_reason: finished.finish_reason,
kv_transfer_params: finished.kv_transfer_params,
})
rust/src/chat/src/output/harmony/tests.rs
+15 -6
View File
@@ -51,8 +51,11 @@ fn decoded_start() -> DecodedTextEvent {
fn finished() -> Finished {
Finished {
prompt_token_count: 0,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}
@@ -112,8 +115,11 @@ fn interrupted_final_message_is_preserved() {
text: "hello".to_string(),
}],
},
prompt_token_count: 0,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
})
@@ -171,8 +177,11 @@ fn interrupted_analysis_message_is_preserved() {
text: "think".to_string(),
}],
},
prompt_token_count: 0,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
})
rust/src/chat/src/output/mod.rs
+3 -4
View File
@@ -5,6 +5,7 @@ use futures::Stream;
use subenum::subenum;
use trait_set::trait_set;
use uuid::Uuid;
use vllm_llm::TokenUsage;
use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
use crate::FinishReason;
@@ -49,8 +50,7 @@ pub(crate) enum AssistantEvent {
ToolCallArgumentsDelta { delta: String },
#[subenum(ContentEvent)]
Done {
prompt_token_count: usize,
output_token_count: usize,
usage: TokenUsage,
finish_reason: FinishReason,
/// Connector-specific KV transfer parameters for disaggregated serving.
kv_transfer_params: Option<serde_json::Value>,
@@ -90,8 +90,7 @@ impl ContentEvent {
}
if let Some(finished) = finished {
events.push(Self::Done {
prompt_token_count: finished.prompt_token_count,
output_token_count: finished.output_token_count,
usage: finished.usage,
finish_reason: finished.finish_reason,
kv_transfer_params: finished.kv_transfer_params,
});
rust/src/chat/src/output/structured.rs
+24 -20
View File
@@ -127,8 +127,7 @@ impl StructuredEventState {
/// Close any open block and emit the terminal `Done` event.
fn finish(
&mut self,
prompt_token_count: usize,
output_token_count: usize,
usage: vllm_llm::TokenUsage,
finish_reason: FinishReason,
kv_transfer_params: Option<serde_json::Value>,
) -> Result<Vec<ChatEvent>> {
@@ -137,8 +136,7 @@ impl StructuredEventState {
self.close_open_tool_call(&mut events);
events.push(ChatEvent::Done {
message: self.message.clone(),
prompt_token_count,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
});
@@ -273,17 +271,11 @@ pub(crate) async fn structured_chat_event_stream(
}
}
AssistantEvent::Done {
prompt_token_count,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
} => {
for next in state.finish(
prompt_token_count,
output_token_count,
finish_reason,
kv_transfer_params,
)? {
for next in state.finish(usage, finish_reason, kv_transfer_params)? {
y.yield_ok(next).await;
}
}
@@ -313,8 +305,11 @@ mod tests {
delta: r#"{"city":"Paris"}"#.to_string(),
}),
Ok(AssistantEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -364,8 +359,11 @@ mod tests {
delta: r#"{"b":2}"#.to_string(),
}),
Ok(AssistantEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -412,8 +410,11 @@ mod tests {
delta: r#"{"city":"Paris"}"#.to_string(),
}),
Ok(AssistantEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -460,8 +461,11 @@ mod tests {
delta: "done".to_string(),
}),
Ok(AssistantEvent::Done {
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
rust/src/chat/src/stream.rs
+13 -10
View File
@@ -14,12 +14,11 @@ use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent};
#[derive(Debug, Clone, PartialEq)]
pub struct CollectedAssistantMessage {
pub message: AssistantMessage,
pub prompt_token_count: usize,
pub prompt_token_ids: Arc<[u32]>,
pub prompt_logprobs: Option<DecodedPromptLogprobs>,
pub logprobs: Option<DecodedLogprobs>,
pub token_ids: Vec<u32>,
pub output_token_count: usize,
pub usage: vllm_llm::TokenUsage,
pub finish_reason: FinishReason,
/// Connector-specific KV transfer parameters for disaggregated serving.
pub kv_transfer_params: Option<serde_json::Value>,
@@ -75,21 +74,19 @@ impl ChatEventStream {
}
ChatEvent::Done {
message: done,
prompt_token_count,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
} => {
return Ok(CollectedAssistantMessage {
message: done,
prompt_token_count,
prompt_token_ids,
prompt_logprobs,
logprobs: (!logprob_positions.is_empty()).then_some(DecodedLogprobs {
positions: logprob_positions,
}),
token_ids,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
});
@@ -190,8 +187,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 2,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 2,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -203,7 +203,6 @@ mod tests {
collected,
CollectedAssistantMessage {
message: Default::default(),
prompt_token_count: 2,
prompt_token_ids: vec![10, 11].into(),
prompt_logprobs: Some(DecodedPromptLogprobs {
first_token_id: 0,
@@ -228,7 +227,11 @@ mod tests {
}],
}),
token_ids: vec![],
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 2,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}
rust/src/chat/tests/chat.rs
+10 -18
View File
@@ -494,12 +494,12 @@ async fn chat_streams_text_events() {
match next_semantic(&mut stream).await {
Some(Ok(ChatEvent::Done {
message,
output_token_count,
usage,
finish_reason,
..
})) => {
assert_eq!(message.text(), "Hi");
assert_eq!(output_token_count, 3);
assert_eq!(usage.output_token_count, 3);
assert_eq!(
finish_reason,
FinishReason::Stop(Some(StopReason::TokenId(b'!' as u32)))
@@ -590,13 +590,9 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
);
match next_semantic(&mut stream).await {
Some(Ok(ChatEvent::Done {
message,
output_token_count,
..
})) => {
Some(Ok(ChatEvent::Done { message, usage, .. })) => {
assert_eq!(message.text(), "你");
assert_eq!(output_token_count, 4);
assert_eq!(usage.output_token_count, 4);
}
other => panic!("unexpected final event: {other:?}"),
}
@@ -681,12 +677,12 @@ async fn chat_stream_flushes_held_text_on_finish() {
match next_semantic(&mut stream).await {
Some(Ok(ChatEvent::Done {
message,
output_token_count,
usage,
finish_reason,
..
})) => {
assert_eq!(message.text(), "ok st");
assert_eq!(output_token_count, 5);
assert_eq!(usage.output_token_count, 5);
assert_eq!(finish_reason, FinishReason::Length);
}
other => panic!("unexpected final event: {other:?}"),
@@ -857,13 +853,9 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
);
match next_semantic(&mut stream).await {
Some(Ok(ChatEvent::Done {
message,
output_token_count,
..
})) => {
Some(Ok(ChatEvent::Done { message, usage, .. })) => {
assert_eq!(message.text(), "Hi!");
assert_eq!(output_token_count, 3);
assert_eq!(usage.output_token_count, 3);
}
other => panic!("unexpected final event: {other:?}"),
}
@@ -1066,11 +1058,11 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
assert_eq!(message.message.text(), "outer");
assert_eq!(message.finish_reason, FinishReason::Length);
assert_eq!(
message.prompt_token_count,
message.usage.prompt_token_count,
"system: You are terse.\nuser: Say hi\nassistant:".len()
);
assert_eq!(
message.output_token_count,
message.usage.output_token_count,
"<think>inner</think>outer".len()
);
rust/src/chat/tests/roundtrip.rs
+10 -4
View File
@@ -508,8 +508,11 @@ fn decoded_completion_stream(
token_ids: Vec::new(),
logprobs: None,
finished: Some(Finished {
prompt_token_count: 0,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -519,8 +522,11 @@ fn decoded_completion_stream(
let last_index = chunks.len() - 1;
for (index, chunk) in chunks.into_iter().enumerate() {
let finished = (index == last_index).then(|| Finished {
prompt_token_count,
output_token_count: completion_body.chars().count(),
usage: vllm_llm::TokenUsage {
prompt_token_count,
output_token_count: completion_body.chars().count(),
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
});
rust/src/cmd/src/cli.rs
+24 -6
View File
@@ -23,8 +23,8 @@ use vllm_engine_core_client::TransportMode;
use vllm_managed_engine::ManagedEngineConfig;
use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
use vllm_server::{
ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
RendererSelection,
ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode,
ParserSelection, RendererSelection,
};
use crate::cli::unsupported::UnsupportedArgs;
@@ -171,6 +171,16 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub enable_log_requests: bool,
/// Include prompt_tokens_details in usage when cached prompt tokens are
/// present.
#[arg(
long,
default_missing_value = "true",
num_args = 0..=1
)]
#[serde(default)]
pub enable_prompt_tokens_details: bool,
/// If specified, API server will add X-Request-Id header to responses.
#[arg(
long,
@@ -248,6 +258,7 @@ impl SharedRuntimeArgs {
) -> Config {
let ready_timeout = self.ready_timeout();
let shutdown_timeout = self.shutdown_timeout();
let api_server_options = self.api_server_options();
Config {
transport_mode: TransportMode::Bootstrapped {
@@ -270,8 +281,7 @@ impl SharedRuntimeArgs {
chat_template: self.chat_template,
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
api_server_options,
api_keys: self.api_key,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
@@ -292,6 +302,7 @@ impl SharedRuntimeArgs {
) -> Config {
let ready_timeout = self.ready_timeout();
let shutdown_timeout = self.shutdown_timeout();
let api_server_options = self.api_server_options();
Config {
transport_mode: TransportMode::HandshakeOwner {
@@ -313,14 +324,21 @@ impl SharedRuntimeArgs {
chat_template: self.chat_template,
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
api_server_options,
api_keys: self.api_key,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
}
}
fn api_server_options(&self) -> ApiServerOptions {
ApiServerOptions {
enable_log_requests: self.enable_log_requests,
enable_prompt_tokens_details: self.enable_prompt_tokens_details,
enable_request_id_headers: self.enable_request_id_headers,
}
}
}
fn default_engine_ready_timeout_secs() -> u64 {
rust/src/cmd/src/cli/tests.rs
+40 -10
View File
@@ -44,6 +44,7 @@ fn serve_args_forward_python_flags_with_separator() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
@@ -143,7 +144,24 @@ fn serve_passes_enable_request_id_headers_into_config() {
panic!("expected serve args");
};
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert!(config.enable_request_id_headers);
assert!(config.api_server_options.enable_request_id_headers);
}
#[test]
fn serve_passes_enable_prompt_tokens_details_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--enable-prompt-tokens-details",
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert!(config.api_server_options.enable_prompt_tokens_details);
}
#[test]
@@ -166,7 +184,7 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() {
panic!("expected frontend args");
};
let config = args.into_config();
assert!(config.enable_request_id_headers);
assert!(config.api_server_options.enable_request_id_headers);
}
#[test]
@@ -342,6 +360,7 @@ fn frontend_args_accept_json() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
@@ -456,7 +475,7 @@ fn frontend_args_json_ignores_unknown_fields() {
}
#[test]
fn frontend_args_json_accepts_noop_fields() {
fn frontend_args_json_sets_prompt_tokens_details_flag() {
let cli = Cli::try_parse_from([
"vllm-rs",
"frontend",
@@ -467,7 +486,7 @@ fn frontend_args_json_accepts_noop_fields() {
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2}"#,
r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2,"enable_prompt_tokens_details":true}"#,
])
.unwrap();
@@ -475,6 +494,7 @@ fn frontend_args_json_accepts_noop_fields() {
panic!("expected frontend args");
};
assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
assert!(args.runtime.enable_prompt_tokens_details);
}
#[test]
@@ -744,6 +764,7 @@ fn serve_args_accept_handshake_aliases() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
@@ -862,8 +883,11 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
chat_template: None,
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_server_options: ApiServerOptions {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
},
api_keys: [],
disable_log_stats: false,
grpc_port: None,
@@ -927,8 +951,11 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
chat_template: None,
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_server_options: ApiServerOptions {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
},
api_keys: [],
disable_log_stats: false,
grpc_port: None,
@@ -1007,8 +1034,11 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
chat_template: None,
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_server_options: ApiServerOptions {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
},
api_keys: [],
disable_log_stats: false,
grpc_port: None,
rust/src/cmd/src/cli/unsupported.rs
-9
View File
@@ -444,15 +444,6 @@ pub struct ServerUnsupportedArgs {
#[arg(long)]
pub max_log_len: Option<Unsupported>,
/// If set to True, enable prompt_tokens_details in usage.
#[arg(
long,
visible_alias = "no-enable-prompt-tokens-details",
default_missing_value = "true",
num_args = 0..=1
)]
pub enable_prompt_tokens_details: Option<Unsupported>,
/// If set to True, enable tracking server_load_metrics in the app state.
#[arg(
long,
rust/src/engine-core-client/src/client.rs
+2 -1
View File
@@ -582,7 +582,8 @@ impl EngineCoreClient {
let results: Vec<T> = self.call_utility(method, args).await?;
if results.iter().all_equal() {
// `engine_count >= 1` is enforced during startup handshake so `results` must be non-empty.
// `engine_count >= 1` is enforced during startup handshake so `results` must be
// non-empty.
Ok(results.into_iter().next().unwrap())
} else {
Err(Error::InconsistentUtilityResults {
rust/src/engine-core-client/src/protocol/utility.rs
+2 -1
View File
@@ -1,5 +1,6 @@
use std::any::type_name;
use std::{fmt, str::FromStr};
use std::fmt;
use std::str::FromStr;
use rmpv::Value;
use serde::{Deserialize, Serialize};
rust/src/llm/src/lib.rs
+1 -1
View File
@@ -10,7 +10,7 @@ mod request_metrics;
pub use error::{Error, Result};
pub use output::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStream,
GenerateOutputStreamExt, GeneratePromptInfo,
GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage,
};
pub use request::GenerateRequest;
pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
rust/src/llm/src/output.rs
+33
View File
@@ -14,6 +14,17 @@ use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
use crate::error::Result;
use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs};
/// Token usage metadata for one request.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenUsage {
/// Number of prompt tokens sent to the engine.
pub prompt_token_count: usize,
/// Number of output tokens generated.
pub output_token_count: usize,
/// Number of prompt tokens served from cache.
pub cached_token_count: usize,
}
/// Final raw token output plus terminal stream metadata.
#[derive(Debug, Clone, PartialEq)]
pub struct CollectedGenerateOutput {
@@ -23,6 +34,7 @@ pub struct CollectedGenerateOutput {
pub token_ids: Vec<u32>,
pub logprobs: Option<Logprobs>,
pub finish_reason: FinishReason,
pub usage: TokenUsage,
/// Connector-specific KV transfer parameters for disaggregated serving.
pub kv_transfer_params: Option<serde_json::Value>,
}
@@ -127,6 +139,8 @@ pub struct GenerateOutput {
pub logprobs: Option<Logprobs>,
/// Terminal finish reason, when this is the final output for the request.
pub finish_reason: Option<FinishReason>,
/// Number of prompt tokens served from cache, when reported by prefill stats.
pub cached_token_count: usize,
/// Connector-specific KV transfer parameters for disaggregated serving.
pub kv_transfer_params: Option<serde_json::Value>,
}
@@ -173,6 +187,7 @@ impl GenerateOutput {
token_ids,
logprobs: None,
finish_reason,
cached_token_count: 0,
kv_transfer_params: None,
}
}
@@ -241,6 +256,11 @@ impl Stream for GenerateOutputStream {
}
let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
let cached_token_count = raw
.prefill_stats
.as_ref()
.map(|stats| stats.num_cached_tokens as usize)
.unwrap_or(0);
let finish_reason = finish_reason_from_engine(raw.finish_reason, raw.stop_reason);
if let Some(finish_reason) = finish_reason.as_ref() {
@@ -253,6 +273,7 @@ impl Stream for GenerateOutputStream {
token_ids: raw.new_token_ids,
logprobs,
finish_reason,
cached_token_count,
kv_transfer_params: raw.kv_transfer_params,
};
@@ -299,9 +320,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
pin_mut!(stream);
let mut prompt_token_ids = None;
let mut prompt_logprobs = None;
let mut cached_token_count = 0;
let mut collected: Option<CollectedGenerateOutput> = None;
while let Some(output) = stream.next().await.transpose()? {
cached_token_count = cached_token_count.max(output.cached_token_count);
if let Some(info) = output.prompt_info {
if prompt_token_ids.is_none() {
prompt_token_ids = Some(info.prompt_token_ids.to_vec());
@@ -328,6 +351,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
token_ids: output.token_ids,
logprobs: output.logprobs,
finish_reason: FinishReason::Error,
usage: TokenUsage {
prompt_token_count: prompt_token_ids.as_ref().map_or(0, Vec::len),
output_token_count: 0,
cached_token_count,
},
kv_transfer_params: None,
});
}
@@ -335,6 +363,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
if let Some(finish_reason) = output.finish_reason {
let mut collected = collected.expect("terminal output must exist");
collected.finish_reason = finish_reason;
collected.usage = TokenUsage {
prompt_token_count: collected.prompt_token_ids.len(),
output_token_count: collected.token_ids.len(),
cached_token_count,
};
collected.kv_transfer_params = output.kv_transfer_params;
return Ok(collected);
}
rust/src/llm/tests/generate.rs
+16 -7
View File
@@ -332,13 +332,21 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
EngineCoreOutputs {
engine_index: 0,
outputs: vec![
request_output_with_logprobs(
&request.request_id,
vec![33],
None,
Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)),
Some(prompt_logprobs()),
),
EngineCoreOutput {
prefill_stats: Some(PrefillStats {
num_prompt_tokens: 2,
num_cached_tokens: 1,
num_local_cached_tokens: 1,
..Default::default()
}),
..request_output_with_logprobs(
&request.request_id,
vec![33],
None,
Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)),
Some(prompt_logprobs()),
)
},
request_output_with_logprobs_and_kv(
&request.request_id,
vec![44],
@@ -373,6 +381,7 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
assert_eq!(collected.prompt_token_ids, vec![11, 22]);
assert_eq!(collected.token_ids, vec![33, 44]);
assert_eq!(collected.finish_reason, FinishReason::stop_eos());
assert_eq!(collected.usage.cached_token_count, 1);
assert_eq!(collected.prompt_logprobs, Some(prompt_logprobs()));
assert_eq!(
collected.logprobs.as_ref().map(|lp| lp.positions.len()),
rust/src/server/examples/external_engine_openai_qwen.rs
+3 -4
View File
@@ -14,8 +14,8 @@ use tokio_util::sync::CancellationToken;
use tracing_subscriber::EnvFilter;
use vllm_engine_core_client::TransportMode;
use vllm_server::{
ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
RendererSelection, serve,
ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode,
ParserSelection, RendererSelection, serve,
};
#[derive(Debug, Parser)]
@@ -68,8 +68,7 @@ async fn main() -> Result<()> {
chat_template: None,
default_chat_template_kwargs: None,
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_server_options: ApiServerOptions::default(),
api_keys: Vec::new(),
disable_log_stats: false,
grpc_port: None,
rust/src/server/src/config.rs
+13 -4
View File
@@ -34,6 +34,17 @@ pub enum CoordinatorMode {
External { address: String },
}
/// HTTP/API-server behavior switches that affect route-layer responses.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
pub struct ApiServerOptions {
/// Log a summary line for each completed request.
pub enable_log_requests: bool,
/// When `true`, include prompt token cache details in response usage.
pub enable_prompt_tokens_details: bool,
/// When `true`, set `X-Request-Id` on every HTTP response.
pub enable_request_id_headers: bool,
}
/// Normalized runtime configuration for the minimal OpenAI-compatible server.
#[derive(Educe, Clone, PartialEq, Eq, Serialize)]
#[educe(Debug)]
@@ -66,10 +77,8 @@ pub struct Config {
pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
/// How to serialize `message.content` for chat-template rendering.
pub chat_template_content_format: ChatTemplateContentFormatOption,
/// Log a summary line for each completed request.
pub enable_log_requests: bool,
/// When `true`, set `X-Request-Id` on every HTTP response.
pub enable_request_id_headers: bool,
/// HTTP/API-server behavior switches.
pub api_server_options: ApiServerOptions,
/// API keys accepted as bearer tokens for guarded routes.
#[serde(skip_serializing)]
#[educe(Debug(method(fmt_redacted_api_keys)))]
rust/src/server/src/grpc/convert.rs
+6 -3
View File
@@ -350,7 +350,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
};
pb::FinishInfo {
num_output_tokens: finished.output_token_count as u32,
num_output_tokens: finished.usage.output_token_count as u32,
finish_reason,
stop_reason,
kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
@@ -590,8 +590,11 @@ mod tests {
fn finished(reason: FinishReason) -> Finished {
Finished {
prompt_token_count: 0,
output_token_count: 0,
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
finish_reason: reason,
kv_transfer_params: None,
}
rust/src/server/src/grpc/mod.rs
+1 -2
View File
@@ -71,8 +71,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl {
);
let finish_info = vllm_text::Finished {
prompt_token_count: collected.prompt_token_ids.len(),
output_token_count: collected.token_ids.len(),
usage: collected.usage,
finish_reason: collected.finish_reason,
kv_transfer_params: collected.kv_transfer_params,
};
rust/src/server/src/lib.rs
+2 -3
View File
@@ -16,7 +16,7 @@ use std::sync::{Arc, OnceLock};
use anyhow::{Context as _, Result};
use axum::Router;
use axum::serve::ListenerExt as _;
pub use config::{Config, CoordinatorMode, HttpListenerMode};
pub use config::{ApiServerOptions, Config, CoordinatorMode, HttpListenerMode};
use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
use tokio_stream::wrappers::TcpListenerStream;
@@ -91,8 +91,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
Ok(Arc::new(
AppState::new(served_model_names, chat)
.with_log_requests(config.enable_log_requests)
.with_request_id_headers(config.enable_request_id_headers)
.with_api_server_options(config.api_server_options)
.with_server_info(ServerInfoSnapshot::from_config(config))
.with_api_keys(config.api_keys.clone()),
))
rust/src/server/src/routes.rs
+1 -1
View File
@@ -100,7 +100,7 @@ fn build_router_with_options(
.route("/server_info", get(server_info::server_info))
}
let enable_request_id_headers = state.enable_request_id_headers;
let enable_request_id_headers = state.api_server_options.enable_request_id_headers;
let enable_api_key_auth = state.has_api_keys();
let mut router = router
.with_state(state.clone())
rust/src/server/src/routes/inference/generate.rs
+52 -21
View File
@@ -19,7 +19,7 @@ use tracing::{error, info, trace};
use tracing_futures::Instrument as _;
use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
use vllm_llm::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage,
};
use self::convert::{ResponseOptions, prepare_generate_request};
@@ -27,6 +27,7 @@ use self::types::{
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
GenerateResponseStreamChoice, GenerateStreamResponse,
};
use crate::config::ApiServerOptions;
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::utils::logprobs::clamp_logprob;
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
@@ -53,7 +54,7 @@ pub async fn generate(
engine_request_id = tracing::field::Empty,
);
let log_request = state.enable_log_requests;
let api_server_options = state.api_server_options;
let stream = prepared.stream;
let raw_stream = match state
.chat
@@ -76,7 +77,7 @@ pub async fn generate(
let chunk_stream = generate_chunk_stream(
raw_stream,
prepared.request_id,
log_request,
api_server_options,
prepared.options,
);
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
@@ -98,7 +99,7 @@ pub async fn generate(
let response = match collect_generate(
collected,
prepared.request_id,
log_request,
api_server_options,
prepared.options,
) {
Ok(response) => response,
@@ -112,7 +113,11 @@ pub async fn generate(
async fn generate_chunk_stream(
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
request_id: String,
log_request: bool,
ApiServerOptions {
enable_log_requests,
enable_prompt_tokens_details,
..
}: ApiServerOptions,
ResponseOptions {
include_usage,
include_continuous_usage,
@@ -123,20 +128,21 @@ async fn generate_chunk_stream(
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
let mut prompt_tokens: Option<u32> = None;
let mut output_tokens = 0_u32;
let mut prompt_tokens = None;
let mut usage = TokenUsage::default();
while let Some(next) = stream.next().await {
match next {
Ok(output) => {
if prompt_tokens.is_none() {
prompt_tokens =
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len());
}
let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
usage.prompt_token_count = prompt_tokens.unwrap_or_default();
usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count);
let token_ids = output.token_ids;
output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
usage.output_token_count = usage.output_token_count.saturating_add(token_ids.len());
let finish_reason = output.finish_reason;
if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
@@ -144,12 +150,12 @@ async fn generate_chunk_stream(
}
if let Some(finish_reason) = finish_reason.as_ref()
&& log_request
&& enable_log_requests
{
info!(
stream = true,
prompt_tokens = usage_prompt_tokens,
output_tokens,
prompt_tokens = usage.prompt_token_count,
output_tokens = usage.output_token_count,
finish_reason = finish_reason.as_str(),
"generate finished"
);
@@ -179,7 +185,7 @@ async fn generate_chunk_stream(
token_ids,
}],
usage: include_continuous_usage
.then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
.then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
})
.await;
}
@@ -197,10 +203,7 @@ async fn generate_chunk_stream(
y.yield_ok(GenerateStreamResponse {
request_id,
choices: Vec::new(),
usage: Some(Usage::from_counts(
prompt_tokens.unwrap_or_default(),
output_tokens,
)),
usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
})
.await;
}
@@ -211,7 +214,10 @@ async fn generate_chunk_stream(
fn collect_generate(
collected: CollectedGenerateOutput,
request_id: String,
log_request: bool,
ApiServerOptions {
enable_log_requests,
..
}: ApiServerOptions,
ResponseOptions {
// Ignored: non-streaming raw generate responses do not include usage.
include_usage: _,
@@ -244,7 +250,7 @@ fn collect_generate(
};
let finish_reason = collected.finish_reason.as_str().to_string();
if log_request {
if enable_log_requests {
info!(
prompt_tokens = collected.prompt_token_ids.len(),
output_tokens = collected.token_ids.len(),
@@ -399,6 +405,7 @@ mod tests {
token_ids: Vec::new(),
logprobs: None,
finish_reason: None,
cached_token_count: 0,
kv_transfer_params: None,
}),
Ok(GenerateOutput {
@@ -410,6 +417,7 @@ mod tests {
token_ids: vec![33],
logprobs: None,
finish_reason: Some(FinishReason::stop_eos()),
cached_token_count: 2,
kv_transfer_params: None,
}),
]);
@@ -417,7 +425,10 @@ mod tests {
let chunks: Vec<_> = generate_chunk_stream(
stream,
"raw-stream".to_string(),
false,
ApiServerOptions {
enable_prompt_tokens_details: true,
..Default::default()
},
ResponseOptions {
include_usage: true,
include_continuous_usage: true,
@@ -433,9 +444,29 @@ mod tests {
chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
2
);
assert_eq!(
chunks[0]
.usage
.as_ref()
.expect("chunk usage")
.prompt_tokens_details
.as_ref()
.map(|details| details.cached_tokens),
Some(2)
);
assert_eq!(
chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
2
);
assert_eq!(
chunks[1]
.usage
.as_ref()
.expect("final usage")
.prompt_tokens_details
.as_ref()
.map(|details| details.cached_tokens),
Some(2)
);
}
}
rust/src/server/src/routes/openai/chat_completions.rs
+77 -35
View File
@@ -24,6 +24,7 @@ use vllm_chat::{
use vllm_engine_core_client::protocol::StopReason;
use self::convert::{ResponseOptions, prepare_chat_request};
use crate::config::ApiServerOptions;
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::chat_completions::types::{
AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
@@ -62,7 +63,7 @@ pub async fn chat_completions(
);
let created = unix_timestamp();
let log_request = state.enable_log_requests;
let api_server_options = state.api_server_options;
let chat_stream =
match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
@@ -82,7 +83,7 @@ pub async fn chat_completions(
prepared.request_id,
prepared.response_model,
created,
log_request,
api_server_options,
prepared.options,
);
let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
@@ -94,7 +95,7 @@ pub async fn chat_completions(
prepared.request_id,
prepared.response_model,
created,
log_request,
api_server_options,
prepared.options,
)
.instrument(request_span.clone())
@@ -113,7 +114,11 @@ async fn collect_chat_completion(
request_id: String,
response_model: String,
created: u64,
log_request: bool,
ApiServerOptions {
enable_log_requests,
enable_prompt_tokens_details,
..
}: ApiServerOptions,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
@@ -133,12 +138,11 @@ async fn collect_chat_completion(
})?;
let CollectedAssistantMessage {
message,
prompt_token_count,
prompt_token_ids,
prompt_logprobs,
logprobs,
token_ids,
output_token_count,
usage,
finish_reason,
kv_transfer_params,
} = collected;
@@ -183,9 +187,9 @@ async fn collect_chat_completion(
} else {
None
};
let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
let usage = Usage::from_token_usage(usage, enable_prompt_tokens_details);
if log_request {
if enable_log_requests {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
@@ -231,7 +235,11 @@ async fn chat_completion_chunk_stream(
request_id: String,
response_model: String,
created: u64,
log_request: bool,
ApiServerOptions {
enable_log_requests,
enable_prompt_tokens_details,
..
}: ApiServerOptions,
ResponseOptions {
include_usage,
requested_logprobs,
@@ -391,17 +399,16 @@ async fn chat_completion_chunk_stream(
debug!("ending current tool call");
}
Ok(ChatEvent::Done {
prompt_token_count,
usage,
finish_reason,
output_token_count,
..
}) => {
if log_request {
if enable_log_requests {
info!(
stream = true,
model = %response_model,
prompt_tokens = prompt_token_count,
output_tokens = output_token_count,
prompt_tokens = usage.prompt_token_count,
output_tokens = usage.output_token_count,
finish_reason = finish_reason.as_str(),
"chat completion finished"
);
@@ -436,7 +443,7 @@ async fn chat_completion_chunk_stream(
&request_id,
&response_model,
created,
Usage::from_counts(prompt_token_count as u32, output_token_count as u32),
Usage::from_token_usage(usage, enable_prompt_tokens_details),
))
.await;
}
@@ -804,7 +811,10 @@ mod tests {
use vllm_engine_core_client::protocol::StopReason;
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::{ResponseOptions, block_delta_chunk, chat_completion_chunk_stream, final_chunk};
use super::{
ApiServerOptions, ResponseOptions, block_delta_chunk, chat_completion_chunk_stream,
final_chunk,
};
#[test]
fn text_chunk_uses_content_only_delta() {
@@ -917,8 +927,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 1,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -929,8 +942,12 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions {
enable_prompt_tokens_details: true,
..Default::default()
},
ResponseOptions {
include_usage: true,
requested_logprobs: true,
include_reasoning: true,
..Default::default()
@@ -942,11 +959,21 @@ mod tests {
.collect::<Result<Vec<_>, _>>()
.expect("stream chunks");
assert_eq!(chunks.len(), 3);
assert_eq!(chunks.len(), 4);
assert_eq!(chunks[1].choices[0].delta.content.as_deref(), Some("hi"));
let logprobs = chunks[1].choices[0].logprobs.as_ref().expect("logprobs");
let content = logprobs.content.as_ref().expect("logprobs content");
assert_eq!(content[0].token, "hi");
assert_eq!(
chunks[3]
.usage
.as_ref()
.expect("usage")
.prompt_tokens_details
.as_ref()
.map(|details| details.cached_tokens),
Some(1)
);
}
#[tokio::test]
@@ -980,8 +1007,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -992,7 +1022,7 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions::default(),
ResponseOptions {
requested_logprobs: true,
include_reasoning: true,
@@ -1032,8 +1062,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 2,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 2,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -1044,7 +1077,7 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions::default(),
ResponseOptions::default(),
)
.collect::<Vec<_>>()
@@ -1110,8 +1143,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 2,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 2,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -1122,7 +1158,7 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions::default(),
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
@@ -1240,8 +1276,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 4,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 4,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -1252,7 +1291,7 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions::default(),
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
@@ -1318,8 +1357,11 @@ mod tests {
}),
Ok(ChatEvent::Done {
message: Default::default(),
prompt_token_count: 1,
output_token_count: 1,
usage: vllm_llm::TokenUsage {
prompt_token_count: 1,
output_token_count: 1,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -1330,7 +1372,7 @@ mod tests {
"chatcmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions::default(),
ResponseOptions {
include_reasoning: true,
..Default::default()
rust/src/server/src/routes/openai/completions.rs
+49 -21
View File
@@ -24,6 +24,7 @@ use super::utils::logprobs::{
text_len,
};
use super::utils::types::Usage;
use crate::config::ApiServerOptions;
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::completions::types::{
CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
@@ -56,7 +57,7 @@ pub async fn completions(
);
let created = unix_timestamp();
let log_request = state.enable_log_requests;
let api_server_options = state.api_server_options;
let text_stream = match state
.chat
.text()
@@ -80,7 +81,7 @@ pub async fn completions(
prepared.request_id,
prepared.response_model,
created,
log_request,
api_server_options,
prepared.options,
);
let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
@@ -92,7 +93,7 @@ pub async fn completions(
prepared.request_id,
prepared.response_model,
created,
log_request,
api_server_options,
prepared.options,
)
.instrument(request_span.clone())
@@ -111,7 +112,11 @@ async fn collect_completion(
request_id: String,
response_model: String,
created: u64,
log_request: bool,
ApiServerOptions {
enable_log_requests,
enable_prompt_tokens_details,
..
}: ApiServerOptions,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
@@ -159,12 +164,9 @@ async fn collect_completion(
Some(prompt) => format!("{prompt}{}", collected.text),
};
let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();
let usage = Usage::from_counts(
collected.prompt_token_ids.len() as u32,
collected.token_ids.len() as u32,
);
let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details);
if log_request {
if enable_log_requests {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
@@ -202,7 +204,11 @@ async fn completion_chunk_stream(
request_id: String,
response_model: String,
created: u64,
log_request: bool,
ApiServerOptions {
enable_log_requests,
enable_prompt_tokens_details,
..
}: ApiServerOptions,
ResponseOptions {
include_usage,
echo,
@@ -275,12 +281,12 @@ async fn completion_chunk_stream(
visible_text_len = visible_text_len.saturating_add(delta_text_len);
if let Some(finished) = finished {
if log_request {
if enable_log_requests {
info!(
stream = true,
model = %response_model,
prompt_tokens = finished.prompt_token_count,
output_tokens = finished.output_token_count,
prompt_tokens = finished.usage.prompt_token_count,
output_tokens = finished.usage.output_token_count,
finish_reason = finished.finish_reason.as_str(),
"completion finished"
);
@@ -298,10 +304,7 @@ async fn completion_chunk_stream(
&request_id,
&response_model,
created,
Usage::from_counts(
finished.prompt_token_count as u32,
finished.output_token_count as u32,
),
Usage::from_token_usage(finished.usage, enable_prompt_tokens_details),
)))
.await;
}
@@ -431,7 +434,9 @@ mod tests {
FinishReason, Finished,
};
use super::{CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk};
use super::{
ApiServerOptions, CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk,
};
#[test]
fn final_chunk_maps_stop_finish_reason() {
@@ -512,8 +517,11 @@ mod tests {
}],
}),
finished: Some(Finished {
prompt_token_count: 5,
output_token_count: 2,
usage: vllm_llm::TokenUsage {
prompt_token_count: 5,
output_token_count: 2,
cached_token_count: 3,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -525,8 +533,12 @@ mod tests {
"cmpl-1".to_string(),
"model".to_string(),
1,
false,
ApiServerOptions {
enable_prompt_tokens_details: true,
..Default::default()
},
ResponseOptions {
include_usage: true,
requested_logprobs: Some(1),
..Default::default()
},
@@ -565,5 +577,21 @@ mod tests {
}
CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),
}
match &chunks[3] {
CompletionSseChunk::Usage(chunk) => {
assert_eq!(
chunk
.usage
.as_ref()
.expect("usage")
.prompt_tokens_details
.as_ref()
.map(|details| details.cached_tokens),
Some(3)
);
}
CompletionSseChunk::Chunk(_) => panic!("expected usage chunk"),
}
}
}
rust/src/server/src/routes/openai/utils/types.rs
+62 -8
View File
@@ -4,6 +4,7 @@ use std::slice;
use llm_multimodal::ImageDetail;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use vllm_llm::TokenUsage;
// ============================================================================
// Constants
@@ -313,29 +314,82 @@ pub enum MessageContent {
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub struct Usage {
pub prompt_tokens: u32,
pub total_tokens: u32,
pub completion_tokens: Option<u32>,
pub prompt_tokens: usize,
pub total_tokens: usize,
pub completion_tokens: Option<usize>,
pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
}
impl Usage {
/// Create a Usage from prompt and completion token counts.
pub fn from_counts(prompt_tokens: u32, completion_tokens: u32) -> Self {
/// Create a Usage with prompt-token cache details.
pub fn from_counts(
prompt_tokens: usize,
completion_tokens: usize,
cached_tokens: Option<usize>,
) -> Self {
Self {
prompt_tokens,
total_tokens: prompt_tokens + completion_tokens,
completion_tokens: Some(completion_tokens),
prompt_tokens_details: None,
prompt_tokens_details: cached_tokens
.filter(|&c| c > 0)
.map(|c| PromptTokenUsageInfo { cached_tokens: c }),
}
}
pub fn from_token_usage(usage: TokenUsage, enable_prompt_tokens_details: bool) -> Self {
Self::from_counts(
usage.prompt_token_count,
usage.output_token_count,
enable_prompt_tokens_details.then_some(usage.cached_token_count),
)
}
}
/// Mirrors the Python vLLM `PromptTokenUsageInfo` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub struct PromptTokenUsageInfo {
pub cached_tokens: Option<u32>,
pub cached_tokens: usize,
}
#[cfg(test)]
mod usage_tests {
use vllm_llm::TokenUsage;
use super::Usage;
#[test]
fn token_usage_hides_prompt_token_details_by_default() {
let usage = Usage::from_token_usage(
TokenUsage {
prompt_token_count: 5,
output_token_count: 2,
cached_token_count: 3,
},
false,
);
assert_eq!(usage.prompt_tokens, 5);
assert_eq!(usage.completion_tokens, Some(2));
assert!(usage.prompt_tokens_details.is_none());
}
#[test]
fn token_usage_includes_prompt_token_details_when_enabled() {
let usage = Usage::from_token_usage(
TokenUsage {
prompt_token_count: 5,
output_token_count: 2,
cached_token_count: 3,
},
true,
);
assert_eq!(
usage.prompt_tokens_details.as_ref().map(|details| details.cached_tokens),
Some(3)
);
}
}
/// OpenAI completions-style logprobs.
rust/src/server/src/routes/tests.rs
+7 -2
View File
@@ -43,6 +43,7 @@ use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
use crate::config::ApiServerOptions;
use crate::state::AppState;
fn request_output(
@@ -787,8 +788,12 @@ async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
)
.await;
let app = build_router(Arc::new(
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
.with_request_id_headers(true),
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_server_options(
ApiServerOptions {
enable_request_id_headers: true,
..Default::default()
},
),
));
(app, engine_task)
}
rust/src/server/src/routes/tokenize/types.rs
+3 -2
View File
@@ -134,11 +134,12 @@ impl Normalizable for DetokenizeRequest {}
#[cfg(test)]
mod tests {
use super::*;
use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
use serde_json::json;
use vllm_chat::ChatTool;
use super::*;
use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
#[test]
fn tokenize_request_converts_openai_tools() {
// The untagged `TokenizeRequest` must resolve a messages+tools body to
rust/src/server/src/state.rs
+7 -15
View File
@@ -9,6 +9,7 @@ use vllm_chat::ChatLlm;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::lora::LoraRequest;
use crate::config::ApiServerOptions;
use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
@@ -27,10 +28,8 @@ pub struct AppState {
served_model_names: Vec<String>,
/// Shared chat facade used by all requests.
pub chat: ChatLlm,
/// Whether to log a summary line for each completed request.
pub enable_log_requests: bool,
/// Whether to set X-Request-Id on every HTTP response.
pub enable_request_id_headers: bool,
/// HTTP/API-server behavior switches.
pub api_server_options: ApiServerOptions,
/// Runtime server information returned by `/server_info`, when available.
server_info: Option<ServerInfoSnapshot>,
/// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
@@ -58,8 +57,7 @@ impl AppState {
Self {
served_model_names,
chat,
enable_log_requests: false,
enable_request_id_headers: false,
api_server_options: ApiServerOptions::default(),
server_info: None,
api_key_hashes: Vec::new(),
server_load: AtomicU64::new(0),
@@ -67,15 +65,9 @@ impl AppState {
}
}
/// Enable per-request completion logging.
pub fn with_log_requests(mut self, enabled: bool) -> Self {
self.enable_log_requests = enabled;
self
}
/// Enable X-Request-Id response headers.
pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
self.enable_request_id_headers = enabled;
/// Set HTTP/API-server behavior switches.
pub fn with_api_server_options(mut self, options: ApiServerOptions) -> Self {
self.api_server_options = options;
self
}
rust/src/text/src/output/decoded.rs
+9 -5
View File
@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
use tracing::{Level, debug, trace};
use vllm_engine_core_client::AbortCause;
use vllm_engine_core_client::protocol::StopReason;
use vllm_llm::{FinishReason, GenerateOutput};
use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
use super::logprobs::{
@@ -40,8 +40,7 @@ impl Default for TextDecodeOptions {
/// Terminal metadata carried on the final [`DecodedTextEvent`].
#[derive(Debug, Clone, PartialEq)]
pub struct Finished {
pub prompt_token_count: usize,
pub output_token_count: usize,
pub usage: TokenUsage,
pub finish_reason: FinishReason,
/// Connector-specific KV transfer parameters for disaggregated serving.
pub kv_transfer_params: Option<serde_json::Value>,
@@ -98,12 +97,14 @@ pub async fn decoded_text_event_stream(
) -> crate::Result<()> {
let mut decoder: Option<Box<dyn IncrementalDecoder>> = None;
let mut prompt_token_count = 0_usize;
let mut cached_token_count = 0_usize;
let mut token_ids = Vec::new();
let mut output_token_count: usize = 0;
let mut logprobs: Option<DecodedLogprobs> = None;
while let Some(next) = raw_stream.next().await {
let output = next?;
cached_token_count = cached_token_count.max(output.cached_token_count);
// If it's the first output, init states and yield `Start` event.
if decoder.is_none() {
@@ -267,8 +268,11 @@ pub async fn decoded_text_event_stream(
token_ids,
logprobs,
finished: Some(Finished {
prompt_token_count,
output_token_count,
usage: TokenUsage {
prompt_token_count,
output_token_count,
cached_token_count,
},
finish_reason: reason,
kv_transfer_params,
}),
rust/src/text/src/output/mod.rs
+13 -4
View File
@@ -23,6 +23,7 @@ pub struct CollectedTextOutput {
pub logprobs: Option<DecodedLogprobs>,
pub token_ids: Vec<u32>,
pub finish_reason: FinishReason,
pub usage: vllm_llm::TokenUsage,
/// Connector-specific KV transfer parameters for disaggregated serving.
pub kv_transfer_params: Option<serde_json::Value>,
}
@@ -74,6 +75,7 @@ impl<T: TextOutputStream> T {
logprobs: delta_logprobs,
token_ids: delta_token_ids,
finish_reason: FinishReason::Error,
usage: vllm_llm::TokenUsage::default(),
kv_transfer_params: None,
})
};
@@ -81,6 +83,7 @@ impl<T: TextOutputStream> T {
if let Some(finished) = finished {
let mut collected = collected.unwrap();
collected.finish_reason = finished.finish_reason;
collected.usage = finished.usage;
collected.kv_transfer_params = finished.kv_transfer_params;
return Ok(collected);
}
@@ -146,8 +149,11 @@ mod tests {
],
}),
finished: Some(Finished {
prompt_token_count: 2,
output_token_count: 2,
usage: vllm_llm::TokenUsage {
prompt_token_count: 2,
output_token_count: 2,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -260,8 +266,11 @@ mod tests {
],
}),
finished: Some(Finished {
prompt_token_count: 2,
output_token_count: 5,
usage: vllm_llm::TokenUsage {
prompt_token_count: 2,
output_token_count: 5,
cached_token_count: 0,
},
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),