From 5d5591d99bb7b2ba695766a992dc328ca5867a4c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 11 Jun 2026 11:50:05 +0800 Subject: [PATCH] [Rust Frontend] Populate `cached_token_count` in responses (#44887) Signed-off-by: Bugen Zhao --- .../examples/external_engine_chat_qwen.rs | 4 +- rust/src/chat/src/event.rs | 7 +- rust/src/chat/src/output/default/reasoning.rs | 17 ++- rust/src/chat/src/output/default/tool.rs | 62 ++++++---- rust/src/chat/src/output/harmony/mod.rs | 3 +- rust/src/chat/src/output/harmony/tests.rs | 21 +++- rust/src/chat/src/output/mod.rs | 7 +- rust/src/chat/src/output/structured.rs | 44 +++---- rust/src/chat/src/stream.rs | 23 ++-- rust/src/chat/tests/chat.rs | 28 ++--- rust/src/chat/tests/roundtrip.rs | 14 ++- rust/src/cmd/src/cli.rs | 30 ++++- rust/src/cmd/src/cli/tests.rs | 50 ++++++-- rust/src/cmd/src/cli/unsupported.rs | 9 -- rust/src/engine-core-client/src/client.rs | 3 +- .../src/protocol/utility.rs | 3 +- rust/src/llm/src/lib.rs | 2 +- rust/src/llm/src/output.rs | 33 ++++++ rust/src/llm/tests/generate.rs | 23 ++-- .../examples/external_engine_openai_qwen.rs | 7 +- rust/src/server/src/config.rs | 17 ++- rust/src/server/src/grpc/convert.rs | 9 +- rust/src/server/src/grpc/mod.rs | 3 +- rust/src/server/src/lib.rs | 5 +- rust/src/server/src/routes.rs | 2 +- .../server/src/routes/inference/generate.rs | 73 ++++++++---- .../src/routes/openai/chat_completions.rs | 112 ++++++++++++------ .../server/src/routes/openai/completions.rs | 70 +++++++---- .../server/src/routes/openai/utils/types.rs | 70 +++++++++-- rust/src/server/src/routes/tests.rs | 9 +- rust/src/server/src/routes/tokenize/types.rs | 5 +- rust/src/server/src/state.rs | 22 ++-- rust/src/text/src/output/decoded.rs | 14 ++- rust/src/text/src/output/mod.rs | 17 ++- 34 files changed, 556 insertions(+), 262 deletions(-) diff --git a/rust/src/chat/examples/external_engine_chat_qwen.rs b/rust/src/chat/examples/external_engine_chat_qwen.rs index d99d672d5e..457dd453d6 100644 --- a/rust/src/chat/examples/external_engine_chat_qwen.rs +++ b/rust/src/chat/examples/external_engine_chat_qwen.rs @@ -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; } diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index 9eb8d35042..d6b5f8f762 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -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, diff --git a/rust/src/chat/src/output/default/reasoning.rs b/rust/src/chat/src/output/default/reasoning.rs index b51ce41961..faa9d7894b 100644 --- a/rust/src/chat/src/output/default/reasoning.rs +++ b/rust/src/chat/src/output/default/reasoning.rs @@ -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, }, diff --git a/rust/src/chat/src/output/default/tool.rs b/rust/src/chat/src/output/default/tool.rs index 8969667530..c216b93f74 100644 --- a/rust/src/chat/src/output/default/tool.rs +++ b/rust/src/chat/src/output/default/tool.rs @@ -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, }), diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 5dc6bc3118..7a043374e5 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -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, }) diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index fe42542b47..91cb52fd0d 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -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, }) diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 6dda8ba0da..d7b73c4e5e 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -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, @@ -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, }); diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index ed6e3a5130..5cbb9f8093 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -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, ) -> Result> { @@ -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, }), diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 8a8dea46e6..fb5c7d3e3f 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -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, pub logprobs: Option, pub token_ids: Vec, - 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, @@ -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, } diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index 7c423561c8..07aa304af0 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -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, "innerouter".len() ); diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 3c8c96ce9f..b3d5d9eae3 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -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, }); diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 624b5da62a..12a85421bd 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -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 { diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 80aa3339db..e351e7e1c8 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -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, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 9a8cdbc279..e9dd5285e5 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -444,15 +444,6 @@ pub struct ServerUnsupportedArgs { #[arg(long)] pub max_log_len: Option, - /// 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, - /// If set to True, enable tracking server_load_metrics in the app state. #[arg( long, diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index f48bf72cd1..7186dfe240 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -582,7 +582,8 @@ impl EngineCoreClient { let results: Vec = 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 { diff --git a/rust/src/engine-core-client/src/protocol/utility.rs b/rust/src/engine-core-client/src/protocol/utility.rs index bfaf2736e0..e15ea6bea0 100644 --- a/rust/src/engine-core-client/src/protocol/utility.rs +++ b/rust/src/engine-core-client/src/protocol/utility.rs @@ -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}; diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index d47935259b..43d46b02f8 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -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}; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 94d9acb3fe..cca7cdca33 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -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, pub logprobs: Option, pub finish_reason: FinishReason, + pub usage: TokenUsage, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, } @@ -127,6 +139,8 @@ pub struct GenerateOutput { pub logprobs: Option, /// Terminal finish reason, when this is the final output for the request. pub finish_reason: Option, + /// 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, } @@ -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> + 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 = 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> + 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> + 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); } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 8b1b98bdc4..18e05063d9 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -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()), diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 8803bd9ea2..510149deea 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -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, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index ac66ea6ce8..aa65dc03c2 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -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>, /// 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)))] diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 0246064b48..0bfe7a63be 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -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, } diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 2f648aa6ce..62ee860766 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -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, }; diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index e2c17cc262..e1257e7f63 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -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> { 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()), )) diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index a942d0a3c9..481c4da461 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -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()) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 5b675f39df..ffbf28048d 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -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>, 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, ) -> Result<(), ApiError> { pin_mut!(stream); - let mut prompt_tokens: Option = 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) + ); } } diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 0a44fabba5..8a2df9b25b 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -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::, _>>() .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::>() @@ -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() diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 8e3c300997..fb0e7bdd87 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -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"), + } } } diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs index 9e0acd04cc..95d16b83b3 100644 --- a/rust/src/server/src/routes/openai/utils/types.rs +++ b/rust/src/server/src/routes/openai/utils/types.rs @@ -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, + pub prompt_tokens: usize, + pub total_tokens: usize, + pub completion_tokens: Option, pub prompt_tokens_details: Option, } 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, + ) -> 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, + 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. diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index a1537d5a1c..68ffe04a3b 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -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) } diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 8067c09f66..9a5977b318 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -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 diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index bcb5f1c6d9..2fee91d457 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -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, /// 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, /// 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 } diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 2ebc6f3853..6452d66b6c 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -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, @@ -98,12 +97,14 @@ pub async fn decoded_text_event_stream( ) -> crate::Result<()> { let mut decoder: Option> = 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 = 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, }), diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 064b820d57..f64d1689f3 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -23,6 +23,7 @@ pub struct CollectedTextOutput { pub logprobs: Option, pub token_ids: Vec, pub finish_reason: FinishReason, + pub usage: vllm_llm::TokenUsage, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, } @@ -74,6 +75,7 @@ impl 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 { 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, }),