@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5d5591d99b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

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

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 11, 2026 at 11:50 UTC 5d5591d99bb7b2ba695766a992dc328ca5867a4c
34 files changed +556 -262
rust/src/chat/examples/external_engine_chat_qwen.rs
+2 -2
@@ -131,13 +131,13 @@ async fn main() -> Result<()> {
131 ChatEvent::LogprobsDelta { .. } => {}
132 ChatEvent::Done {
133 message,
134 - output_token_count,
134 + usage,
135 finish_reason: reason,
136 ..
137 } => {
138 final_reasoning = message.reasoning().unwrap_or_default();
139 final_text = message.text();
140 - final_output_token_count = output_token_count;
140 + final_output_token_count = usage.output_token_count;
141 finish_reason = Some(reason);
142 break;
143 }
rust/src/chat/src/event.rs
+2 -5
@@ -2,6 +2,7 @@ use std::ops::Deref;
2 use std::sync::Arc;
3
4 use serde::{Deserialize, Serialize};
5 +use vllm_llm::TokenUsage;
6 use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
7
8 use crate::FinishReason;
@@ -197,11 +198,7 @@ pub enum ChatEvent {
198 /// metadata.
199 Done {
200 message: AssistantMessage,
200 - /// Number of prompt tokens actually sent to the engine after chat
201 - /// template rendering and tokenization.
202 - prompt_token_count: usize,
203 - /// Number of output tokens generated.
204 - output_token_count: usize,
201 + usage: TokenUsage,
202 finish_reason: FinishReason,
203 /// Connector-specific KV transfer parameters for disaggregated serving.
204 kv_transfer_params: Option<serde_json::Value>,
rust/src/chat/src/output/default/reasoning.rs
+11 -6
@@ -178,8 +178,7 @@ pub(crate) async fn reasoning_event_stream(
178 y.yield_ok(next).await;
179 }
180 y.yield_ok(ContentEvent::Done {
181 - prompt_token_count: finished.prompt_token_count,
182 - output_token_count: finished.output_token_count,
181 + usage: finished.usage,
182 finish_reason: finished.finish_reason,
183 kv_transfer_params: finished.kv_transfer_params,
184 })
@@ -289,8 +288,11 @@ mod tests {
288 token_ids: vec![],
289 logprobs: None,
290 finished: Some(vllm_text::Finished {
292 - prompt_token_count: 3,
293 - output_token_count: 0,
291 + usage: vllm_llm::TokenUsage {
292 + prompt_token_count: 3,
293 + output_token_count: 0,
294 + cached_token_count: 0,
295 + },
296 finish_reason: FinishReason::stop_eos(),
297 kv_transfer_params: None,
298 }),
@@ -322,8 +324,11 @@ mod tests {
324 delta: "def".to_string(),
325 },
326 ContentEvent::Done {
325 - prompt_token_count: 3,
326 - output_token_count: 0,
327 + usage: vllm_llm::TokenUsage {
328 + prompt_token_count: 3,
329 + output_token_count: 0,
330 + cached_token_count: 0,
331 + },
332 finish_reason: FinishReason::stop_eos(),
333 kv_transfer_params: None,
334 },
rust/src/chat/src/output/default/tool.rs
+42 -20
@@ -240,8 +240,7 @@ pub(crate) async fn tool_event_stream(
240 .await;
241 }
242 ContentEvent::Done {
243 - prompt_token_count,
244 - output_token_count,
243 + usage,
244 finish_reason,
245 kv_transfer_params,
246 } => {
@@ -250,8 +249,7 @@ pub(crate) async fn tool_event_stream(
249 }
250
251 y.yield_ok(AssistantEvent::Done {
253 - prompt_token_count,
254 - output_token_count,
252 + usage,
253 finish_reason,
254 kv_transfer_params,
255 })
@@ -465,8 +463,11 @@ mod tests {
463 })
464 })
465 .chain(std::iter::once(Ok(ContentEvent::Done {
468 - prompt_token_count: 1,
469 - output_token_count: 1,
466 + usage: vllm_llm::TokenUsage {
467 + prompt_token_count: 1,
468 + output_token_count: 1,
469 + cached_token_count: 0,
470 + },
471 finish_reason: FinishReason::stop_eos(),
472 kv_transfer_params: None,
473 })));
@@ -506,8 +507,11 @@ mod tests {
507 delta: "ignored".to_string(),
508 }),
509 Ok(ContentEvent::Done {
509 - prompt_token_count: 1,
510 - output_token_count: 1,
510 + usage: vllm_llm::TokenUsage {
511 + prompt_token_count: 1,
512 + output_token_count: 1,
513 + cached_token_count: 0,
514 + },
515 finish_reason: FinishReason::stop_eos(),
516 kv_transfer_params: None,
517 }),
@@ -659,8 +663,11 @@ mod tests {
663 delta: "def".to_string(),
664 }),
665 Ok(ContentEvent::Done {
662 - prompt_token_count: 3,
663 - output_token_count: 0,
666 + usage: vllm_llm::TokenUsage {
667 + prompt_token_count: 3,
668 + output_token_count: 0,
669 + cached_token_count: 0,
670 + },
671 finish_reason: FinishReason::stop_eos(),
672 kv_transfer_params: None,
673 }),
@@ -697,8 +704,11 @@ mod tests {
704 delta: "def".to_string(),
705 },
706 AssistantEvent::Done {
700 - prompt_token_count: 3,
701 - output_token_count: 0,
707 + usage: vllm_llm::TokenUsage {
708 + prompt_token_count: 3,
709 + output_token_count: 0,
710 + cached_token_count: 0,
711 + },
712 finish_reason: FinishReason::stop_eos(),
713 kv_transfer_params: None,
714 },
@@ -739,8 +749,11 @@ mod tests {
749 token_ids: vec![],
750 }),
751 Ok(ContentEvent::Done {
742 - prompt_token_count: 1,
743 - output_token_count: 0,
752 + usage: vllm_llm::TokenUsage {
753 + prompt_token_count: 1,
754 + output_token_count: 0,
755 + cached_token_count: 0,
756 + },
757 finish_reason: FinishReason::stop_eos(),
758 kv_transfer_params: None,
759 }),
@@ -779,8 +792,11 @@ mod tests {
792 token_ids: vec![],
793 },
794 AssistantEvent::Done {
782 - prompt_token_count: 1,
783 - output_token_count: 0,
795 + usage: vllm_llm::TokenUsage {
796 + prompt_token_count: 1,
797 + output_token_count: 0,
798 + cached_token_count: 0,
799 + },
800 finish_reason: FinishReason::stop_eos(),
801 kv_transfer_params: None,
802 },
@@ -796,8 +812,11 @@ mod tests {
812 delta: "ignored".to_string(),
813 }),
814 Ok(ContentEvent::Done {
799 - prompt_token_count: 1,
800 - output_token_count: 1,
815 + usage: vllm_llm::TokenUsage {
816 + prompt_token_count: 1,
817 + output_token_count: 1,
818 + cached_token_count: 0,
819 + },
820 finish_reason: FinishReason::stop_eos(),
821 kv_transfer_params: None,
822 }),
@@ -901,8 +920,11 @@ mod tests {
920 delta: "ignored".to_string(),
921 }),
922 Ok(ContentEvent::Done {
904 - prompt_token_count: 1,
905 - output_token_count: 1,
923 + usage: vllm_llm::TokenUsage {
924 + prompt_token_count: 1,
925 + output_token_count: 1,
926 + cached_token_count: 0,
927 + },
928 finish_reason: FinishReason::stop_eos(),
929 kv_transfer_params: None,
930 }),
rust/src/chat/src/output/harmony/mod.rs
+1 -2
@@ -366,8 +366,7 @@ async fn harmony_assistant_event_stream(
366
367 if let Some(finished) = finished {
368 y.yield_ok(AssistantEvent::Done {
369 - prompt_token_count: finished.prompt_token_count,
370 - output_token_count: finished.output_token_count,
369 + usage: finished.usage,
370 finish_reason: finished.finish_reason,
371 kv_transfer_params: finished.kv_transfer_params,
372 })
rust/src/chat/src/output/harmony/tests.rs
+15 -6
@@ -51,8 +51,11 @@ fn decoded_start() -> DecodedTextEvent {
51
52 fn finished() -> Finished {
53 Finished {
54 - prompt_token_count: 0,
55 - output_token_count: 0,
54 + usage: vllm_llm::TokenUsage {
55 + prompt_token_count: 0,
56 + output_token_count: 0,
57 + cached_token_count: 0,
58 + },
59 finish_reason: FinishReason::stop_eos(),
60 kv_transfer_params: None,
61 }
@@ -112,8 +115,11 @@ fn interrupted_final_message_is_preserved() {
115 text: "hello".to_string(),
116 }],
117 },
115 - prompt_token_count: 0,
116 - output_token_count: 0,
118 + usage: vllm_llm::TokenUsage {
119 + prompt_token_count: 0,
120 + output_token_count: 0,
121 + cached_token_count: 0,
122 + },
123 finish_reason: FinishReason::stop_eos(),
124 kv_transfer_params: None,
125 })
@@ -171,8 +177,11 @@ fn interrupted_analysis_message_is_preserved() {
177 text: "think".to_string(),
178 }],
179 },
174 - prompt_token_count: 0,
175 - output_token_count: 0,
180 + usage: vllm_llm::TokenUsage {
181 + prompt_token_count: 0,
182 + output_token_count: 0,
183 + cached_token_count: 0,
184 + },
185 finish_reason: FinishReason::stop_eos(),
186 kv_transfer_params: None,
187 })
rust/src/chat/src/output/mod.rs
+3 -4
@@ -5,6 +5,7 @@ use futures::Stream;
5 use subenum::subenum;
6 use trait_set::trait_set;
7 use uuid::Uuid;
8 +use vllm_llm::TokenUsage;
9 use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
10
11 use crate::FinishReason;
@@ -49,8 +50,7 @@ pub(crate) enum AssistantEvent {
50 ToolCallArgumentsDelta { delta: String },
51 #[subenum(ContentEvent)]
52 Done {
52 - prompt_token_count: usize,
53 - output_token_count: usize,
53 + usage: TokenUsage,
54 finish_reason: FinishReason,
55 /// Connector-specific KV transfer parameters for disaggregated serving.
56 kv_transfer_params: Option<serde_json::Value>,
@@ -90,8 +90,7 @@ impl ContentEvent {
90 }
91 if let Some(finished) = finished {
92 events.push(Self::Done {
93 - prompt_token_count: finished.prompt_token_count,
94 - output_token_count: finished.output_token_count,
93 + usage: finished.usage,
94 finish_reason: finished.finish_reason,
95 kv_transfer_params: finished.kv_transfer_params,
96 });
rust/src/chat/src/output/structured.rs
+24 -20
@@ -127,8 +127,7 @@ impl StructuredEventState {
127 /// Close any open block and emit the terminal `Done` event.
128 fn finish(
129 &mut self,
130 - prompt_token_count: usize,
131 - output_token_count: usize,
130 + usage: vllm_llm::TokenUsage,
131 finish_reason: FinishReason,
132 kv_transfer_params: Option<serde_json::Value>,
133 ) -> Result<Vec<ChatEvent>> {
@@ -137,8 +136,7 @@ impl StructuredEventState {
136 self.close_open_tool_call(&mut events);
137 events.push(ChatEvent::Done {
138 message: self.message.clone(),
140 - prompt_token_count,
141 - output_token_count,
139 + usage,
140 finish_reason,
141 kv_transfer_params,
142 });
@@ -273,17 +271,11 @@ pub(crate) async fn structured_chat_event_stream(
271 }
272 }
273 AssistantEvent::Done {
276 - prompt_token_count,
277 - output_token_count,
274 + usage,
275 finish_reason,
276 kv_transfer_params,
277 } => {
281 - for next in state.finish(
282 - prompt_token_count,
283 - output_token_count,
284 - finish_reason,
285 - kv_transfer_params,
286 - )? {
278 + for next in state.finish(usage, finish_reason, kv_transfer_params)? {
279 y.yield_ok(next).await;
280 }
281 }
@@ -313,8 +305,11 @@ mod tests {
305 delta: r#"{"city":"Paris"}"#.to_string(),
306 }),
307 Ok(AssistantEvent::Done {
316 - prompt_token_count: 1,
317 - output_token_count: 1,
308 + usage: vllm_llm::TokenUsage {
309 + prompt_token_count: 1,
310 + output_token_count: 1,
311 + cached_token_count: 0,
312 + },
313 finish_reason: FinishReason::stop_eos(),
314 kv_transfer_params: None,
315 }),
@@ -364,8 +359,11 @@ mod tests {
359 delta: r#"{"b":2}"#.to_string(),
360 }),
361 Ok(AssistantEvent::Done {
367 - prompt_token_count: 1,
368 - output_token_count: 1,
362 + usage: vllm_llm::TokenUsage {
363 + prompt_token_count: 1,
364 + output_token_count: 1,
365 + cached_token_count: 0,
366 + },
367 finish_reason: FinishReason::stop_eos(),
368 kv_transfer_params: None,
369 }),
@@ -412,8 +410,11 @@ mod tests {
410 delta: r#"{"city":"Paris"}"#.to_string(),
411 }),
412 Ok(AssistantEvent::Done {
415 - prompt_token_count: 1,
416 - output_token_count: 1,
413 + usage: vllm_llm::TokenUsage {
414 + prompt_token_count: 1,
415 + output_token_count: 1,
416 + cached_token_count: 0,
417 + },
418 finish_reason: FinishReason::stop_eos(),
419 kv_transfer_params: None,
420 }),
@@ -460,8 +461,11 @@ mod tests {
461 delta: "done".to_string(),
462 }),
463 Ok(AssistantEvent::Done {
463 - prompt_token_count: 1,
464 - output_token_count: 1,
464 + usage: vllm_llm::TokenUsage {
465 + prompt_token_count: 1,
466 + output_token_count: 1,
467 + cached_token_count: 0,
468 + },
469 finish_reason: FinishReason::stop_eos(),
470 kv_transfer_params: None,
471 }),
rust/src/chat/src/stream.rs
+13 -10
@@ -14,12 +14,11 @@ use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent};
14 #[derive(Debug, Clone, PartialEq)]
15 pub struct CollectedAssistantMessage {
16 pub message: AssistantMessage,
17 - pub prompt_token_count: usize,
17 pub prompt_token_ids: Arc<[u32]>,
18 pub prompt_logprobs: Option<DecodedPromptLogprobs>,
19 pub logprobs: Option<DecodedLogprobs>,
20 pub token_ids: Vec<u32>,
22 - pub output_token_count: usize,
21 + pub usage: vllm_llm::TokenUsage,
22 pub finish_reason: FinishReason,
23 /// Connector-specific KV transfer parameters for disaggregated serving.
24 pub kv_transfer_params: Option<serde_json::Value>,
@@ -75,21 +74,19 @@ impl ChatEventStream {
74 }
75 ChatEvent::Done {
76 message: done,
78 - prompt_token_count,
79 - output_token_count,
77 + usage,
78 finish_reason,
79 kv_transfer_params,
80 } => {
81 return Ok(CollectedAssistantMessage {
82 message: done,
85 - prompt_token_count,
83 prompt_token_ids,
84 prompt_logprobs,
85 logprobs: (!logprob_positions.is_empty()).then_some(DecodedLogprobs {
86 positions: logprob_positions,
87 }),
88 token_ids,
92 - output_token_count,
89 + usage,
90 finish_reason,
91 kv_transfer_params,
92 });
@@ -190,8 +187,11 @@ mod tests {
187 }),
188 Ok(ChatEvent::Done {
189 message: Default::default(),
193 - prompt_token_count: 2,
194 - output_token_count: 1,
190 + usage: vllm_llm::TokenUsage {
191 + prompt_token_count: 2,
192 + output_token_count: 1,
193 + cached_token_count: 0,
194 + },
195 finish_reason: FinishReason::stop_eos(),
196 kv_transfer_params: None,
197 }),
@@ -203,7 +203,6 @@ mod tests {
203 collected,
204 CollectedAssistantMessage {
205 message: Default::default(),
206 - prompt_token_count: 2,
206 prompt_token_ids: vec![10, 11].into(),
207 prompt_logprobs: Some(DecodedPromptLogprobs {
208 first_token_id: 0,
@@ -228,7 +227,11 @@ mod tests {
227 }],
228 }),
229 token_ids: vec![],
231 - output_token_count: 1,
230 + usage: vllm_llm::TokenUsage {
231 + prompt_token_count: 2,
232 + output_token_count: 1,
233 + cached_token_count: 0,
234 + },
235 finish_reason: FinishReason::stop_eos(),
236 kv_transfer_params: None,
237 }
rust/src/chat/tests/chat.rs
+10 -18
@@ -494,12 +494,12 @@ async fn chat_streams_text_events() {
494 match next_semantic(&mut stream).await {
495 Some(Ok(ChatEvent::Done {
496 message,
497 - output_token_count,
497 + usage,
498 finish_reason,
499 ..
500 })) => {
501 assert_eq!(message.text(), "Hi");
502 - assert_eq!(output_token_count, 3);
502 + assert_eq!(usage.output_token_count, 3);
503 assert_eq!(
504 finish_reason,
505 FinishReason::Stop(Some(StopReason::TokenId(b'!' as u32)))
@@ -590,13 +590,9 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
590 );
591
592 match next_semantic(&mut stream).await {
593 - Some(Ok(ChatEvent::Done {
594 - message,
595 - output_token_count,
596 - ..
597 - })) => {
593 + Some(Ok(ChatEvent::Done { message, usage, .. })) => {
594 assert_eq!(message.text(), "你");
599 - assert_eq!(output_token_count, 4);
595 + assert_eq!(usage.output_token_count, 4);
596 }
597 other => panic!("unexpected final event: {other:?}"),
598 }
@@ -681,12 +677,12 @@ async fn chat_stream_flushes_held_text_on_finish() {
677 match next_semantic(&mut stream).await {
678 Some(Ok(ChatEvent::Done {
679 message,
684 - output_token_count,
680 + usage,
681 finish_reason,
682 ..
683 })) => {
684 assert_eq!(message.text(), "ok st");
689 - assert_eq!(output_token_count, 5);
685 + assert_eq!(usage.output_token_count, 5);
686 assert_eq!(finish_reason, FinishReason::Length);
687 }
688 other => panic!("unexpected final event: {other:?}"),
@@ -857,13 +853,9 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
853 );
854
855 match next_semantic(&mut stream).await {
860 - Some(Ok(ChatEvent::Done {
861 - message,
862 - output_token_count,
863 - ..
864 - })) => {
856 + Some(Ok(ChatEvent::Done { message, usage, .. })) => {
857 assert_eq!(message.text(), "Hi!");
866 - assert_eq!(output_token_count, 3);
858 + assert_eq!(usage.output_token_count, 3);
859 }
860 other => panic!("unexpected final event: {other:?}"),
861 }
@@ -1066,11 +1058,11 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
1058 assert_eq!(message.message.text(), "outer");
1059 assert_eq!(message.finish_reason, FinishReason::Length);
1060 assert_eq!(
1069 - message.prompt_token_count,
1061 + message.usage.prompt_token_count,
1062 "system: You are terse.\nuser: Say hi\nassistant:".len()
1063 );
1064 assert_eq!(
1073 - message.output_token_count,
1065 + message.usage.output_token_count,
1066 "<think>inner</think>outer".len()
1067 );
1068
rust/src/chat/tests/roundtrip.rs
+10 -4
@@ -508,8 +508,11 @@ fn decoded_completion_stream(
508 token_ids: Vec::new(),
509 logprobs: None,
510 finished: Some(Finished {
511 - prompt_token_count: 0,
512 - output_token_count: 0,
511 + usage: vllm_llm::TokenUsage {
512 + prompt_token_count: 0,
513 + output_token_count: 0,
514 + cached_token_count: 0,
515 + },
516 finish_reason: FinishReason::stop_eos(),
517 kv_transfer_params: None,
518 }),
@@ -519,8 +522,11 @@ fn decoded_completion_stream(
522 let last_index = chunks.len() - 1;
523 for (index, chunk) in chunks.into_iter().enumerate() {
524 let finished = (index == last_index).then(|| Finished {
522 - prompt_token_count,
523 - output_token_count: completion_body.chars().count(),
525 + usage: vllm_llm::TokenUsage {
526 + prompt_token_count,
527 + output_token_count: completion_body.chars().count(),
528 + cached_token_count: 0,
529 + },
530 finish_reason: FinishReason::stop_eos(),
531 kv_transfer_params: None,
532 });
rust/src/cmd/src/cli.rs
+24 -6
@@ -23,8 +23,8 @@ use vllm_engine_core_client::TransportMode;
23 use vllm_managed_engine::ManagedEngineConfig;
24 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
25 use vllm_server::{
26 - ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
27 - RendererSelection,
26 + ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode,
27 + ParserSelection, RendererSelection,
28 };
29
30 use crate::cli::unsupported::UnsupportedArgs;
@@ -171,6 +171,16 @@ pub struct SharedRuntimeArgs {
171 #[serde(default)]
172 pub enable_log_requests: bool,
173
174 + /// Include prompt_tokens_details in usage when cached prompt tokens are
175 + /// present.
176 + #[arg(
177 + long,
178 + default_missing_value = "true",
179 + num_args = 0..=1
180 + )]
181 + #[serde(default)]
182 + pub enable_prompt_tokens_details: bool,
183 +
184 /// If specified, API server will add X-Request-Id header to responses.
185 #[arg(
186 long,
@@ -248,6 +258,7 @@ impl SharedRuntimeArgs {
258 ) -> Config {
259 let ready_timeout = self.ready_timeout();
260 let shutdown_timeout = self.shutdown_timeout();
261 + let api_server_options = self.api_server_options();
262
263 Config {
264 transport_mode: TransportMode::Bootstrapped {
@@ -270,8 +281,7 @@ impl SharedRuntimeArgs {
281 chat_template: self.chat_template,
282 default_chat_template_kwargs: self.default_chat_template_kwargs,
283 chat_template_content_format: self.chat_template_content_format,
273 - enable_log_requests: self.enable_log_requests,
274 - enable_request_id_headers: self.enable_request_id_headers,
284 + api_server_options,
285 api_keys: self.api_key,
286 disable_log_stats: self.disable_log_stats,
287 grpc_port: self.grpc_port,
@@ -292,6 +302,7 @@ impl SharedRuntimeArgs {
302 ) -> Config {
303 let ready_timeout = self.ready_timeout();
304 let shutdown_timeout = self.shutdown_timeout();
305 + let api_server_options = self.api_server_options();
306
307 Config {
308 transport_mode: TransportMode::HandshakeOwner {
@@ -313,14 +324,21 @@ impl SharedRuntimeArgs {
324 chat_template: self.chat_template,
325 default_chat_template_kwargs: self.default_chat_template_kwargs,
326 chat_template_content_format: self.chat_template_content_format,
316 - enable_log_requests: self.enable_log_requests,
317 - enable_request_id_headers: self.enable_request_id_headers,
327 + api_server_options,
328 api_keys: self.api_key,
329 disable_log_stats: self.disable_log_stats,
330 grpc_port: self.grpc_port,
331 shutdown_timeout,
332 }
333 }
334 +
335 + fn api_server_options(&self) -> ApiServerOptions {
336 + ApiServerOptions {
337 + enable_log_requests: self.enable_log_requests,
338 + enable_prompt_tokens_details: self.enable_prompt_tokens_details,
339 + enable_request_id_headers: self.enable_request_id_headers,
340 + }
341 + }
342 }
343
344 fn default_engine_ready_timeout_secs() -> u64 {
rust/src/cmd/src/cli/tests.rs
+40 -10
@@ -44,6 +44,7 @@ fn serve_args_forward_python_flags_with_separator() {
44 default_chat_template_kwargs: None,
45 chat_template_content_format: Auto,
46 enable_log_requests: false,
47 + enable_prompt_tokens_details: false,
48 enable_request_id_headers: false,
49 disable_log_stats: false,
50 served_model_name: [],
@@ -143,7 +144,24 @@ fn serve_passes_enable_request_id_headers_into_config() {
144 panic!("expected serve args");
145 };
146 let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
146 - assert!(config.enable_request_id_headers);
147 + assert!(config.api_server_options.enable_request_id_headers);
148 +}
149 +
150 +#[test]
151 +fn serve_passes_enable_prompt_tokens_details_into_config() {
152 + let cli = Cli::try_parse_from([
153 + "vllm-rs",
154 + "serve",
155 + "Qwen/Qwen3-0.6B",
156 + "--enable-prompt-tokens-details",
157 + ])
158 + .unwrap();
159 +
160 + let Command::Serve(args) = cli.command else {
161 + panic!("expected serve args");
162 + };
163 + let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
164 + assert!(config.api_server_options.enable_prompt_tokens_details);
165 }
166
167 #[test]
@@ -166,7 +184,7 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() {
184 panic!("expected frontend args");
185 };
186 let config = args.into_config();
169 - assert!(config.enable_request_id_headers);
187 + assert!(config.api_server_options.enable_request_id_headers);
188 }
189
190 #[test]
@@ -342,6 +360,7 @@ fn frontend_args_accept_json() {
360 default_chat_template_kwargs: None,
361 chat_template_content_format: Auto,
362 enable_log_requests: false,
363 + enable_prompt_tokens_details: false,
364 enable_request_id_headers: false,
365 disable_log_stats: false,
366 served_model_name: [],
@@ -456,7 +475,7 @@ fn frontend_args_json_ignores_unknown_fields() {
475 }
476
477 #[test]
459 -fn frontend_args_json_accepts_noop_fields() {
478 +fn frontend_args_json_sets_prompt_tokens_details_flag() {
479 let cli = Cli::try_parse_from([
480 "vllm-rs",
481 "frontend",
@@ -467,7 +486,7 @@ fn frontend_args_json_accepts_noop_fields() {
486 "--output-address",
487 "ipc:///tmp/output.sock",
488 "--args-json",
470 - r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2}"#,
489 + r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2,"enable_prompt_tokens_details":true}"#,
490 ])
491 .unwrap();
492
@@ -475,6 +494,7 @@ fn frontend_args_json_accepts_noop_fields() {
494 panic!("expected frontend args");
495 };
496 assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
497 + assert!(args.runtime.enable_prompt_tokens_details);
498 }
499
500 #[test]
@@ -744,6 +764,7 @@ fn serve_args_accept_handshake_aliases() {
764 default_chat_template_kwargs: None,
765 chat_template_content_format: Auto,
766 enable_log_requests: false,
767 + enable_prompt_tokens_details: false,
768 enable_request_id_headers: false,
769 disable_log_stats: false,
770 served_model_name: [],
@@ -862,8 +883,11 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
883 chat_template: None,
884 default_chat_template_kwargs: None,
885 chat_template_content_format: Auto,
865 - enable_log_requests: false,
866 - enable_request_id_headers: false,
886 + api_server_options: ApiServerOptions {
887 + enable_log_requests: false,
888 + enable_prompt_tokens_details: false,
889 + enable_request_id_headers: false,
890 + },
891 api_keys: [],
892 disable_log_stats: false,
893 grpc_port: None,
@@ -927,8 +951,11 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
951 chat_template: None,
952 default_chat_template_kwargs: None,
953 chat_template_content_format: Auto,
930 - enable_log_requests: false,
931 - enable_request_id_headers: false,
954 + api_server_options: ApiServerOptions {
955 + enable_log_requests: false,
956 + enable_prompt_tokens_details: false,
957 + enable_request_id_headers: false,
958 + },
959 api_keys: [],
960 disable_log_stats: false,
961 grpc_port: None,
@@ -1007,8 +1034,11 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1034 chat_template: None,
1035 default_chat_template_kwargs: None,
1036 chat_template_content_format: Auto,
1010 - enable_log_requests: false,
1011 - enable_request_id_headers: false,
1037 + api_server_options: ApiServerOptions {
1038 + enable_log_requests: false,
1039 + enable_prompt_tokens_details: false,
1040 + enable_request_id_headers: false,
1041 + },
1042 api_keys: [],
1043 disable_log_stats: false,
1044 grpc_port: None,
rust/src/cmd/src/cli/unsupported.rs
-9
@@ -444,15 +444,6 @@ pub struct ServerUnsupportedArgs {
444 #[arg(long)]
445 pub max_log_len: Option<Unsupported>,
446
447 - /// If set to True, enable prompt_tokens_details in usage.
448 - #[arg(
449 - long,
450 - visible_alias = "no-enable-prompt-tokens-details",
451 - default_missing_value = "true",
452 - num_args = 0..=1
453 - )]
454 - pub enable_prompt_tokens_details: Option<Unsupported>,
455 -
447 /// If set to True, enable tracking server_load_metrics in the app state.
448 #[arg(
449 long,
rust/src/engine-core-client/src/client.rs
+2 -1
@@ -582,7 +582,8 @@ impl EngineCoreClient {
582 let results: Vec<T> = self.call_utility(method, args).await?;
583
584 if results.iter().all_equal() {
585 - // `engine_count >= 1` is enforced during startup handshake so `results` must be non-empty.
585 + // `engine_count >= 1` is enforced during startup handshake so `results` must be
586 + // non-empty.
587 Ok(results.into_iter().next().unwrap())
588 } else {
589 Err(Error::InconsistentUtilityResults {
rust/src/engine-core-client/src/protocol/utility.rs
+2 -1
@@ -1,5 +1,6 @@
1 use std::any::type_name;
2 -use std::{fmt, str::FromStr};
2 +use std::fmt;
3 +use std::str::FromStr;
4
5 use rmpv::Value;
6 use serde::{Deserialize, Serialize};
rust/src/llm/src/lib.rs
+1 -1
@@ -10,7 +10,7 @@ mod request_metrics;
10 pub use error::{Error, Result};
11 pub use output::{
12 CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStream,
13 - GenerateOutputStreamExt, GeneratePromptInfo,
13 + GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage,
14 };
15 pub use request::GenerateRequest;
16 pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
rust/src/llm/src/output.rs
+33
@@ -14,6 +14,17 @@ use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
14 use crate::error::Result;
15 use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs};
16
17 +/// Token usage metadata for one request.
18 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
19 +pub struct TokenUsage {
20 + /// Number of prompt tokens sent to the engine.
21 + pub prompt_token_count: usize,
22 + /// Number of output tokens generated.
23 + pub output_token_count: usize,
24 + /// Number of prompt tokens served from cache.
25 + pub cached_token_count: usize,
26 +}
27 +
28 /// Final raw token output plus terminal stream metadata.
29 #[derive(Debug, Clone, PartialEq)]
30 pub struct CollectedGenerateOutput {
@@ -23,6 +34,7 @@ pub struct CollectedGenerateOutput {
34 pub token_ids: Vec<u32>,
35 pub logprobs: Option<Logprobs>,
36 pub finish_reason: FinishReason,
37 + pub usage: TokenUsage,
38 /// Connector-specific KV transfer parameters for disaggregated serving.
39 pub kv_transfer_params: Option<serde_json::Value>,
40 }
@@ -127,6 +139,8 @@ pub struct GenerateOutput {
139 pub logprobs: Option<Logprobs>,
140 /// Terminal finish reason, when this is the final output for the request.
141 pub finish_reason: Option<FinishReason>,
142 + /// Number of prompt tokens served from cache, when reported by prefill stats.
143 + pub cached_token_count: usize,
144 /// Connector-specific KV transfer parameters for disaggregated serving.
145 pub kv_transfer_params: Option<serde_json::Value>,
146 }
@@ -173,6 +187,7 @@ impl GenerateOutput {
187 token_ids,
188 logprobs: None,
189 finish_reason,
190 + cached_token_count: 0,
191 kv_transfer_params: None,
192 }
193 }
@@ -241,6 +256,11 @@ impl Stream for GenerateOutputStream {
256 }
257
258 let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
259 + let cached_token_count = raw
260 + .prefill_stats
261 + .as_ref()
262 + .map(|stats| stats.num_cached_tokens as usize)
263 + .unwrap_or(0);
264
265 let finish_reason = finish_reason_from_engine(raw.finish_reason, raw.stop_reason);
266 if let Some(finish_reason) = finish_reason.as_ref() {
@@ -253,6 +273,7 @@ impl Stream for GenerateOutputStream {
273 token_ids: raw.new_token_ids,
274 logprobs,
275 finish_reason,
276 + cached_token_count,
277 kv_transfer_params: raw.kv_transfer_params,
278 };
279
@@ -299,9 +320,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
320 pin_mut!(stream);
321 let mut prompt_token_ids = None;
322 let mut prompt_logprobs = None;
323 + let mut cached_token_count = 0;
324 let mut collected: Option<CollectedGenerateOutput> = None;
325
326 while let Some(output) = stream.next().await.transpose()? {
327 + cached_token_count = cached_token_count.max(output.cached_token_count);
328 if let Some(info) = output.prompt_info {
329 if prompt_token_ids.is_none() {
330 prompt_token_ids = Some(info.prompt_token_ids.to_vec());
@@ -328,6 +351,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
351 token_ids: output.token_ids,
352 logprobs: output.logprobs,
353 finish_reason: FinishReason::Error,
354 + usage: TokenUsage {
355 + prompt_token_count: prompt_token_ids.as_ref().map_or(0, Vec::len),
356 + output_token_count: 0,
357 + cached_token_count,
358 + },
359 kv_transfer_params: None,
360 });
361 }
@@ -335,6 +363,11 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
363 if let Some(finish_reason) = output.finish_reason {
364 let mut collected = collected.expect("terminal output must exist");
365 collected.finish_reason = finish_reason;
366 + collected.usage = TokenUsage {
367 + prompt_token_count: collected.prompt_token_ids.len(),
368 + output_token_count: collected.token_ids.len(),
369 + cached_token_count,
370 + };
371 collected.kv_transfer_params = output.kv_transfer_params;
372 return Ok(collected);
373 }
rust/src/llm/tests/generate.rs
+16 -7
@@ -332,13 +332,21 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
332 EngineCoreOutputs {
333 engine_index: 0,
334 outputs: vec![
335 - request_output_with_logprobs(
336 - &request.request_id,
337 - vec![33],
338 - None,
339 - Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)),
340 - Some(prompt_logprobs()),
341 - ),
335 + EngineCoreOutput {
336 + prefill_stats: Some(PrefillStats {
337 + num_prompt_tokens: 2,
338 + num_cached_tokens: 1,
339 + num_local_cached_tokens: 1,
340 + ..Default::default()
341 + }),
342 + ..request_output_with_logprobs(
343 + &request.request_id,
344 + vec![33],
345 + None,
346 + Some(logprobs_for_position(33, -0.1, 1, 99, -0.2)),
347 + Some(prompt_logprobs()),
348 + )
349 + },
350 request_output_with_logprobs_and_kv(
351 &request.request_id,
352 vec![44],
@@ -373,6 +381,7 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
381 assert_eq!(collected.prompt_token_ids, vec![11, 22]);
382 assert_eq!(collected.token_ids, vec![33, 44]);
383 assert_eq!(collected.finish_reason, FinishReason::stop_eos());
384 + assert_eq!(collected.usage.cached_token_count, 1);
385 assert_eq!(collected.prompt_logprobs, Some(prompt_logprobs()));
386 assert_eq!(
387 collected.logprobs.as_ref().map(|lp| lp.positions.len()),
rust/src/server/examples/external_engine_openai_qwen.rs
+3 -4
@@ -14,8 +14,8 @@ use tokio_util::sync::CancellationToken;
14 use tracing_subscriber::EnvFilter;
15 use vllm_engine_core_client::TransportMode;
16 use vllm_server::{
17 - ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode, ParserSelection,
18 - RendererSelection, serve,
17 + ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode,
18 + ParserSelection, RendererSelection, serve,
19 };
20
21 #[derive(Debug, Parser)]
@@ -68,8 +68,7 @@ async fn main() -> Result<()> {
68 chat_template: None,
69 default_chat_template_kwargs: None,
70 chat_template_content_format: ChatTemplateContentFormatOption::Auto,
71 - enable_log_requests: false,
72 - enable_request_id_headers: false,
71 + api_server_options: ApiServerOptions::default(),
72 api_keys: Vec::new(),
73 disable_log_stats: false,
74 grpc_port: None,
rust/src/server/src/config.rs
+13 -4
@@ -34,6 +34,17 @@ pub enum CoordinatorMode {
34 External { address: String },
35 }
36
37 +/// HTTP/API-server behavior switches that affect route-layer responses.
38 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
39 +pub struct ApiServerOptions {
40 + /// Log a summary line for each completed request.
41 + pub enable_log_requests: bool,
42 + /// When `true`, include prompt token cache details in response usage.
43 + pub enable_prompt_tokens_details: bool,
44 + /// When `true`, set `X-Request-Id` on every HTTP response.
45 + pub enable_request_id_headers: bool,
46 +}
47 +
48 /// Normalized runtime configuration for the minimal OpenAI-compatible server.
49 #[derive(Educe, Clone, PartialEq, Eq, Serialize)]
50 #[educe(Debug)]
@@ -66,10 +77,8 @@ pub struct Config {
77 pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
78 /// How to serialize `message.content` for chat-template rendering.
79 pub chat_template_content_format: ChatTemplateContentFormatOption,
69 - /// Log a summary line for each completed request.
70 - pub enable_log_requests: bool,
71 - /// When `true`, set `X-Request-Id` on every HTTP response.
72 - pub enable_request_id_headers: bool,
80 + /// HTTP/API-server behavior switches.
81 + pub api_server_options: ApiServerOptions,
82 /// API keys accepted as bearer tokens for guarded routes.
83 #[serde(skip_serializing)]
84 #[educe(Debug(method(fmt_redacted_api_keys)))]
rust/src/server/src/grpc/convert.rs
+6 -3
@@ -350,7 +350,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
350 };
351
352 pb::FinishInfo {
353 - num_output_tokens: finished.output_token_count as u32,
353 + num_output_tokens: finished.usage.output_token_count as u32,
354 finish_reason,
355 stop_reason,
356 kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
@@ -590,8 +590,11 @@ mod tests {
590
591 fn finished(reason: FinishReason) -> Finished {
592 Finished {
593 - prompt_token_count: 0,
594 - output_token_count: 0,
593 + usage: vllm_llm::TokenUsage {
594 + prompt_token_count: 0,
595 + output_token_count: 0,
596 + cached_token_count: 0,
597 + },
598 finish_reason: reason,
599 kv_transfer_params: None,
600 }
rust/src/server/src/grpc/mod.rs
+1 -2
@@ -71,8 +71,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl {
71 );
72
73 let finish_info = vllm_text::Finished {
74 - prompt_token_count: collected.prompt_token_ids.len(),
75 - output_token_count: collected.token_ids.len(),
74 + usage: collected.usage,
75 finish_reason: collected.finish_reason,
76 kv_transfer_params: collected.kv_transfer_params,
77 };
rust/src/server/src/lib.rs
+2 -3
@@ -16,7 +16,7 @@ use std::sync::{Arc, OnceLock};
16 use anyhow::{Context as _, Result};
17 use axum::Router;
18 use axum::serve::ListenerExt as _;
19 -pub use config::{Config, CoordinatorMode, HttpListenerMode};
19 +pub use config::{ApiServerOptions, Config, CoordinatorMode, HttpListenerMode};
20 use tokio::net::TcpListener;
21 use tokio::time::{Instant, sleep_until};
22 use tokio_stream::wrappers::TcpListenerStream;
@@ -91,8 +91,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
91
92 Ok(Arc::new(
93 AppState::new(served_model_names, chat)
94 - .with_log_requests(config.enable_log_requests)
95 - .with_request_id_headers(config.enable_request_id_headers)
94 + .with_api_server_options(config.api_server_options)
95 .with_server_info(ServerInfoSnapshot::from_config(config))
96 .with_api_keys(config.api_keys.clone()),
97 ))
rust/src/server/src/routes.rs
+1 -1
@@ -100,7 +100,7 @@ fn build_router_with_options(
100 .route("/server_info", get(server_info::server_info))
101 }
102
103 - let enable_request_id_headers = state.enable_request_id_headers;
103 + let enable_request_id_headers = state.api_server_options.enable_request_id_headers;
104 let enable_api_key_auth = state.has_api_keys();
105 let mut router = router
106 .with_state(state.clone())
rust/src/server/src/routes/inference/generate.rs
+52 -21
@@ -19,7 +19,7 @@ use tracing::{error, info, trace};
19 use tracing_futures::Instrument as _;
20 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
21 use vllm_llm::{
22 - CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
22 + CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage,
23 };
24
25 use self::convert::{ResponseOptions, prepare_generate_request};
@@ -27,6 +27,7 @@ use self::types::{
27 GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
28 GenerateResponseStreamChoice, GenerateStreamResponse,
29 };
30 +use crate::config::ApiServerOptions;
31 use crate::error::{ApiError, bail_server_error, server_error};
32 use crate::routes::openai::utils::logprobs::clamp_logprob;
33 use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
@@ -53,7 +54,7 @@ pub async fn generate(
54 engine_request_id = tracing::field::Empty,
55 );
56
56 - let log_request = state.enable_log_requests;
57 + let api_server_options = state.api_server_options;
58 let stream = prepared.stream;
59 let raw_stream = match state
60 .chat
@@ -76,7 +77,7 @@ pub async fn generate(
77 let chunk_stream = generate_chunk_stream(
78 raw_stream,
79 prepared.request_id,
79 - log_request,
80 + api_server_options,
81 prepared.options,
82 );
83 let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
@@ -98,7 +99,7 @@ pub async fn generate(
99 let response = match collect_generate(
100 collected,
101 prepared.request_id,
101 - log_request,
102 + api_server_options,
103 prepared.options,
104 ) {
105 Ok(response) => response,
@@ -112,7 +113,11 @@ pub async fn generate(
113 async fn generate_chunk_stream(
114 stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
115 request_id: String,
115 - log_request: bool,
116 + ApiServerOptions {
117 + enable_log_requests,
118 + enable_prompt_tokens_details,
119 + ..
120 + }: ApiServerOptions,
121 ResponseOptions {
122 include_usage,
123 include_continuous_usage,
@@ -123,20 +128,21 @@ async fn generate_chunk_stream(
128 mut y: TryYielder<GenerateStreamResponse, ApiError>,
129 ) -> Result<(), ApiError> {
130 pin_mut!(stream);
126 - let mut prompt_tokens: Option<u32> = None;
127 - let mut output_tokens = 0_u32;
131 + let mut prompt_tokens = None;
132 + let mut usage = TokenUsage::default();
133
134 while let Some(next) = stream.next().await {
135 match next {
136 Ok(output) => {
137 if prompt_tokens.is_none() {
138 prompt_tokens =
134 - output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
139 + output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len());
140 }
136 - let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
141 + usage.prompt_token_count = prompt_tokens.unwrap_or_default();
142 + usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count);
143
144 let token_ids = output.token_ids;
139 - output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
145 + usage.output_token_count = usage.output_token_count.saturating_add(token_ids.len());
146 let finish_reason = output.finish_reason;
147
148 if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
@@ -144,12 +150,12 @@ async fn generate_chunk_stream(
150 }
151
152 if let Some(finish_reason) = finish_reason.as_ref()
147 - && log_request
153 + && enable_log_requests
154 {
155 info!(
156 stream = true,
151 - prompt_tokens = usage_prompt_tokens,
152 - output_tokens,
157 + prompt_tokens = usage.prompt_token_count,
158 + output_tokens = usage.output_token_count,
159 finish_reason = finish_reason.as_str(),
160 "generate finished"
161 );
@@ -179,7 +185,7 @@ async fn generate_chunk_stream(
185 token_ids,
186 }],
187 usage: include_continuous_usage
182 - .then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
188 + .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
189 })
190 .await;
191 }
@@ -197,10 +203,7 @@ async fn generate_chunk_stream(
203 y.yield_ok(GenerateStreamResponse {
204 request_id,
205 choices: Vec::new(),
200 - usage: Some(Usage::from_counts(
201 - prompt_tokens.unwrap_or_default(),
202 - output_tokens,
203 - )),
206 + usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
207 })
208 .await;
209 }
@@ -211,7 +214,10 @@ async fn generate_chunk_stream(
214 fn collect_generate(
215 collected: CollectedGenerateOutput,
216 request_id: String,
214 - log_request: bool,
217 + ApiServerOptions {
218 + enable_log_requests,
219 + ..
220 + }: ApiServerOptions,
221 ResponseOptions {
222 // Ignored: non-streaming raw generate responses do not include usage.
223 include_usage: _,
@@ -244,7 +250,7 @@ fn collect_generate(
250 };
251 let finish_reason = collected.finish_reason.as_str().to_string();
252
247 - if log_request {
253 + if enable_log_requests {
254 info!(
255 prompt_tokens = collected.prompt_token_ids.len(),
256 output_tokens = collected.token_ids.len(),
@@ -399,6 +405,7 @@ mod tests {
405 token_ids: Vec::new(),
406 logprobs: None,
407 finish_reason: None,
408 + cached_token_count: 0,
409 kv_transfer_params: None,
410 }),
411 Ok(GenerateOutput {
@@ -410,6 +417,7 @@ mod tests {
417 token_ids: vec![33],
418 logprobs: None,
419 finish_reason: Some(FinishReason::stop_eos()),
420 + cached_token_count: 2,
421 kv_transfer_params: None,
422 }),
423 ]);
@@ -417,7 +425,10 @@ mod tests {
425 let chunks: Vec<_> = generate_chunk_stream(
426 stream,
427 "raw-stream".to_string(),
420 - false,
428 + ApiServerOptions {
429 + enable_prompt_tokens_details: true,
430 + ..Default::default()
431 + },
432 ResponseOptions {
433 include_usage: true,
434 include_continuous_usage: true,
@@ -433,9 +444,29 @@ mod tests {
444 chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
445 2
446 );
447 + assert_eq!(
448 + chunks[0]
449 + .usage
450 + .as_ref()
451 + .expect("chunk usage")
452 + .prompt_tokens_details
453 + .as_ref()
454 + .map(|details| details.cached_tokens),
455 + Some(2)
456 + );
457 assert_eq!(
458 chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
459 2
460 );
461 + assert_eq!(
462 + chunks[1]
463 + .usage
464 + .as_ref()
465 + .expect("final usage")
466 + .prompt_tokens_details
467 + .as_ref()
468 + .map(|details| details.cached_tokens),
469 + Some(2)
470 + );
471 }
472 }
rust/src/server/src/routes/openai/chat_completions.rs
+77 -35
@@ -24,6 +24,7 @@ use vllm_chat::{
24 use vllm_engine_core_client::protocol::StopReason;
25
26 use self::convert::{ResponseOptions, prepare_chat_request};
27 +use crate::config::ApiServerOptions;
28 use crate::error::{ApiError, bail_server_error, server_error};
29 use crate::routes::openai::chat_completions::types::{
30 AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
@@ -62,7 +63,7 @@ pub async fn chat_completions(
63 );
64
65 let created = unix_timestamp();
65 - let log_request = state.enable_log_requests;
66 + let api_server_options = state.api_server_options;
67
68 let chat_stream =
69 match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
@@ -82,7 +83,7 @@ pub async fn chat_completions(
83 prepared.request_id,
84 prepared.response_model,
85 created,
85 - log_request,
86 + api_server_options,
87 prepared.options,
88 );
89 let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
@@ -94,7 +95,7 @@ pub async fn chat_completions(
95 prepared.request_id,
96 prepared.response_model,
97 created,
97 - log_request,
98 + api_server_options,
99 prepared.options,
100 )
101 .instrument(request_span.clone())
@@ -113,7 +114,11 @@ async fn collect_chat_completion(
114 request_id: String,
115 response_model: String,
116 created: u64,
116 - log_request: bool,
117 + ApiServerOptions {
118 + enable_log_requests,
119 + enable_prompt_tokens_details,
120 + ..
121 + }: ApiServerOptions,
122 ResponseOptions {
123 // Ignored: non-streaming responses always include usage.
124 include_usage: _,
@@ -133,12 +138,11 @@ async fn collect_chat_completion(
138 })?;
139 let CollectedAssistantMessage {
140 message,
136 - prompt_token_count,
141 prompt_token_ids,
142 prompt_logprobs,
143 logprobs,
144 token_ids,
141 - output_token_count,
145 + usage,
146 finish_reason,
147 kv_transfer_params,
148 } = collected;
@@ -183,9 +187,9 @@ async fn collect_chat_completion(
187 } else {
188 None
189 };
186 - let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
190 + let usage = Usage::from_token_usage(usage, enable_prompt_tokens_details);
191
188 - if log_request {
192 + if enable_log_requests {
193 info!(
194 model = %response_model,
195 prompt_tokens = usage.prompt_tokens,
@@ -231,7 +235,11 @@ async fn chat_completion_chunk_stream(
235 request_id: String,
236 response_model: String,
237 created: u64,
234 - log_request: bool,
238 + ApiServerOptions {
239 + enable_log_requests,
240 + enable_prompt_tokens_details,
241 + ..
242 + }: ApiServerOptions,
243 ResponseOptions {
244 include_usage,
245 requested_logprobs,
@@ -391,17 +399,16 @@ async fn chat_completion_chunk_stream(
399 debug!("ending current tool call");
400 }
401 Ok(ChatEvent::Done {
394 - prompt_token_count,
402 + usage,
403 finish_reason,
396 - output_token_count,
404 ..
405 }) => {
399 - if log_request {
406 + if enable_log_requests {
407 info!(
408 stream = true,
409 model = %response_model,
403 - prompt_tokens = prompt_token_count,
404 - output_tokens = output_token_count,
410 + prompt_tokens = usage.prompt_token_count,
411 + output_tokens = usage.output_token_count,
412 finish_reason = finish_reason.as_str(),
413 "chat completion finished"
414 );
@@ -436,7 +443,7 @@ async fn chat_completion_chunk_stream(
443 &request_id,
444 &response_model,
445 created,
439 - Usage::from_counts(prompt_token_count as u32, output_token_count as u32),
446 + Usage::from_token_usage(usage, enable_prompt_tokens_details),
447 ))
448 .await;
449 }
@@ -804,7 +811,10 @@ mod tests {
811 use vllm_engine_core_client::protocol::StopReason;
812 use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
813
807 - use super::{ResponseOptions, block_delta_chunk, chat_completion_chunk_stream, final_chunk};
814 + use super::{
815 + ApiServerOptions, ResponseOptions, block_delta_chunk, chat_completion_chunk_stream,
816 + final_chunk,
817 + };
818
819 #[test]
820 fn text_chunk_uses_content_only_delta() {
@@ -917,8 +927,11 @@ mod tests {
927 }),
928 Ok(ChatEvent::Done {
929 message: Default::default(),
920 - prompt_token_count: 1,
921 - output_token_count: 1,
930 + usage: vllm_llm::TokenUsage {
931 + prompt_token_count: 1,
932 + output_token_count: 1,
933 + cached_token_count: 1,
934 + },
935 finish_reason: FinishReason::stop_eos(),
936 kv_transfer_params: None,
937 }),
@@ -929,8 +942,12 @@ mod tests {
942 "chatcmpl-1".to_string(),
943 "model".to_string(),
944 1,
932 - false,
945 + ApiServerOptions {
946 + enable_prompt_tokens_details: true,
947 + ..Default::default()
948 + },
949 ResponseOptions {
950 + include_usage: true,
951 requested_logprobs: true,
952 include_reasoning: true,
953 ..Default::default()
@@ -942,11 +959,21 @@ mod tests {
959 .collect::<Result<Vec<_>, _>>()
960 .expect("stream chunks");
961
945 - assert_eq!(chunks.len(), 3);
962 + assert_eq!(chunks.len(), 4);
963 assert_eq!(chunks[1].choices[0].delta.content.as_deref(), Some("hi"));
964 let logprobs = chunks[1].choices[0].logprobs.as_ref().expect("logprobs");
965 let content = logprobs.content.as_ref().expect("logprobs content");
966 assert_eq!(content[0].token, "hi");
967 + assert_eq!(
968 + chunks[3]
969 + .usage
970 + .as_ref()
971 + .expect("usage")
972 + .prompt_tokens_details
973 + .as_ref()
974 + .map(|details| details.cached_tokens),
975 + Some(1)
976 + );
977 }
978
979 #[tokio::test]
@@ -980,8 +1007,11 @@ mod tests {
1007 }),
1008 Ok(ChatEvent::Done {
1009 message: Default::default(),
983 - prompt_token_count: 1,
984 - output_token_count: 1,
1010 + usage: vllm_llm::TokenUsage {
1011 + prompt_token_count: 1,
1012 + output_token_count: 1,
1013 + cached_token_count: 0,
1014 + },
1015 finish_reason: FinishReason::stop_eos(),
1016 kv_transfer_params: None,
1017 }),
@@ -992,7 +1022,7 @@ mod tests {
1022 "chatcmpl-1".to_string(),
1023 "model".to_string(),
1024 1,
995 - false,
1025 + ApiServerOptions::default(),
1026 ResponseOptions {
1027 requested_logprobs: true,
1028 include_reasoning: true,
@@ -1032,8 +1062,11 @@ mod tests {
1062 }),
1063 Ok(ChatEvent::Done {
1064 message: Default::default(),
1035 - prompt_token_count: 1,
1036 - output_token_count: 2,
1065 + usage: vllm_llm::TokenUsage {
1066 + prompt_token_count: 1,
1067 + output_token_count: 2,
1068 + cached_token_count: 0,
1069 + },
1070 finish_reason: FinishReason::stop_eos(),
1071 kv_transfer_params: None,
1072 }),
@@ -1044,7 +1077,7 @@ mod tests {
1077 "chatcmpl-1".to_string(),
1078 "model".to_string(),
1079 1,
1047 - false,
1080 + ApiServerOptions::default(),
1081 ResponseOptions::default(),
1082 )
1083 .collect::<Vec<_>>()
@@ -1110,8 +1143,11 @@ mod tests {
1143 }),
1144 Ok(ChatEvent::Done {
1145 message: Default::default(),
1113 - prompt_token_count: 1,
1114 - output_token_count: 2,
1146 + usage: vllm_llm::TokenUsage {
1147 + prompt_token_count: 1,
1148 + output_token_count: 2,
1149 + cached_token_count: 0,
1150 + },
1151 finish_reason: FinishReason::stop_eos(),
1152 kv_transfer_params: None,
1153 }),
@@ -1122,7 +1158,7 @@ mod tests {
1158 "chatcmpl-1".to_string(),
1159 "model".to_string(),
1160 1,
1125 - false,
1161 + ApiServerOptions::default(),
1162 ResponseOptions {
1163 requested_logprobs: true,
1164 return_token_ids: true,
@@ -1240,8 +1276,11 @@ mod tests {
1276 }),
1277 Ok(ChatEvent::Done {
1278 message: Default::default(),
1243 - prompt_token_count: 1,
1244 - output_token_count: 4,
1279 + usage: vllm_llm::TokenUsage {
1280 + prompt_token_count: 1,
1281 + output_token_count: 4,
1282 + cached_token_count: 0,
1283 + },
1284 finish_reason: FinishReason::stop_eos(),
1285 kv_transfer_params: None,
1286 }),
@@ -1252,7 +1291,7 @@ mod tests {
1291 "chatcmpl-1".to_string(),
1292 "model".to_string(),
1293 1,
1255 - false,
1294 + ApiServerOptions::default(),
1295 ResponseOptions {
1296 requested_logprobs: true,
1297 return_token_ids: true,
@@ -1318,8 +1357,11 @@ mod tests {
1357 }),
1358 Ok(ChatEvent::Done {
1359 message: Default::default(),
1321 - prompt_token_count: 1,
1322 - output_token_count: 1,
1360 + usage: vllm_llm::TokenUsage {
1361 + prompt_token_count: 1,
1362 + output_token_count: 1,
1363 + cached_token_count: 0,
1364 + },
1365 finish_reason: FinishReason::stop_eos(),
1366 kv_transfer_params: None,
1367 }),
@@ -1330,7 +1372,7 @@ mod tests {
1372 "chatcmpl-1".to_string(),
1373 "model".to_string(),
1374 1,
1333 - false,
1375 + ApiServerOptions::default(),
1376 ResponseOptions {
1377 include_reasoning: true,
1378 ..Default::default()
rust/src/server/src/routes/openai/completions.rs
+49 -21
@@ -24,6 +24,7 @@ use super::utils::logprobs::{
24 text_len,
25 };
26 use super::utils::types::Usage;
27 +use crate::config::ApiServerOptions;
28 use crate::error::{ApiError, bail_server_error, server_error};
29 use crate::routes::openai::completions::types::{
30 CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
@@ -56,7 +57,7 @@ pub async fn completions(
57 );
58
59 let created = unix_timestamp();
59 - let log_request = state.enable_log_requests;
60 + let api_server_options = state.api_server_options;
61 let text_stream = match state
62 .chat
63 .text()
@@ -80,7 +81,7 @@ pub async fn completions(
81 prepared.request_id,
82 prepared.response_model,
83 created,
83 - log_request,
84 + api_server_options,
85 prepared.options,
86 );
87 let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
@@ -92,7 +93,7 @@ pub async fn completions(
93 prepared.request_id,
94 prepared.response_model,
95 created,
95 - log_request,
96 + api_server_options,
97 prepared.options,
98 )
99 .instrument(request_span.clone())
@@ -111,7 +112,11 @@ async fn collect_completion(
112 request_id: String,
113 response_model: String,
114 created: u64,
114 - log_request: bool,
115 + ApiServerOptions {
116 + enable_log_requests,
117 + enable_prompt_tokens_details,
118 + ..
119 + }: ApiServerOptions,
120 ResponseOptions {
121 // Ignored: non-streaming responses always include usage.
122 include_usage: _,
@@ -159,12 +164,9 @@ async fn collect_completion(
164 Some(prompt) => format!("{prompt}{}", collected.text),
165 };
166 let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();
162 - let usage = Usage::from_counts(
163 - collected.prompt_token_ids.len() as u32,
164 - collected.token_ids.len() as u32,
165 - );
167 + let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details);
168
167 - if log_request {
169 + if enable_log_requests {
170 info!(
171 model = %response_model,
172 prompt_tokens = usage.prompt_tokens,
@@ -202,7 +204,11 @@ async fn completion_chunk_stream(
204 request_id: String,
205 response_model: String,
206 created: u64,
205 - log_request: bool,
207 + ApiServerOptions {
208 + enable_log_requests,
209 + enable_prompt_tokens_details,
210 + ..
211 + }: ApiServerOptions,
212 ResponseOptions {
213 include_usage,
214 echo,
@@ -275,12 +281,12 @@ async fn completion_chunk_stream(
281 visible_text_len = visible_text_len.saturating_add(delta_text_len);
282
283 if let Some(finished) = finished {
278 - if log_request {
284 + if enable_log_requests {
285 info!(
286 stream = true,
287 model = %response_model,
282 - prompt_tokens = finished.prompt_token_count,
283 - output_tokens = finished.output_token_count,
288 + prompt_tokens = finished.usage.prompt_token_count,
289 + output_tokens = finished.usage.output_token_count,
290 finish_reason = finished.finish_reason.as_str(),
291 "completion finished"
292 );
@@ -298,10 +304,7 @@ async fn completion_chunk_stream(
304 &request_id,
305 &response_model,
306 created,
301 - Usage::from_counts(
302 - finished.prompt_token_count as u32,
303 - finished.output_token_count as u32,
304 - ),
307 + Usage::from_token_usage(finished.usage, enable_prompt_tokens_details),
308 )))
309 .await;
310 }
@@ -431,7 +434,9 @@ mod tests {
434 FinishReason, Finished,
435 };
436
434 - use super::{CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk};
437 + use super::{
438 + ApiServerOptions, CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk,
439 + };
440
441 #[test]
442 fn final_chunk_maps_stop_finish_reason() {
@@ -512,8 +517,11 @@ mod tests {
517 }],
518 }),
519 finished: Some(Finished {
515 - prompt_token_count: 5,
516 - output_token_count: 2,
520 + usage: vllm_llm::TokenUsage {
521 + prompt_token_count: 5,
522 + output_token_count: 2,
523 + cached_token_count: 3,
524 + },
525 finish_reason: FinishReason::stop_eos(),
526 kv_transfer_params: None,
527 }),
@@ -525,8 +533,12 @@ mod tests {
533 "cmpl-1".to_string(),
534 "model".to_string(),
535 1,
528 - false,
536 + ApiServerOptions {
537 + enable_prompt_tokens_details: true,
538 + ..Default::default()
539 + },
540 ResponseOptions {
541 + include_usage: true,
542 requested_logprobs: Some(1),
543 ..Default::default()
544 },
@@ -565,5 +577,21 @@ mod tests {
577 }
578 CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),
579 }
580 +
581 + match &chunks[3] {
582 + CompletionSseChunk::Usage(chunk) => {
583 + assert_eq!(
584 + chunk
585 + .usage
586 + .as_ref()
587 + .expect("usage")
588 + .prompt_tokens_details
589 + .as_ref()
590 + .map(|details| details.cached_tokens),
591 + Some(3)
592 + );
593 + }
594 + CompletionSseChunk::Chunk(_) => panic!("expected usage chunk"),
595 + }
596 }
597 }
rust/src/server/src/routes/openai/utils/types.rs
+62 -8
@@ -4,6 +4,7 @@ use std::slice;
4 use llm_multimodal::ImageDetail;
5 use serde::{Deserialize, Serialize};
6 use serde_json::Value;
7 +use vllm_llm::TokenUsage;
8
9 // ============================================================================
10 // Constants
@@ -313,29 +314,82 @@ pub enum MessageContent {
314 #[serde_with::skip_serializing_none]
315 #[derive(Debug, Clone, Serialize)]
316 pub struct Usage {
316 - pub prompt_tokens: u32,
317 - pub total_tokens: u32,
318 - pub completion_tokens: Option<u32>,
317 + pub prompt_tokens: usize,
318 + pub total_tokens: usize,
319 + pub completion_tokens: Option<usize>,
320 pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
321 }
322
323 impl Usage {
323 - /// Create a Usage from prompt and completion token counts.
324 - pub fn from_counts(prompt_tokens: u32, completion_tokens: u32) -> Self {
324 + /// Create a Usage with prompt-token cache details.
325 + pub fn from_counts(
326 + prompt_tokens: usize,
327 + completion_tokens: usize,
328 + cached_tokens: Option<usize>,
329 + ) -> Self {
330 Self {
331 prompt_tokens,
332 total_tokens: prompt_tokens + completion_tokens,
333 completion_tokens: Some(completion_tokens),
329 - prompt_tokens_details: None,
334 + prompt_tokens_details: cached_tokens
335 + .filter(|&c| c > 0)
336 + .map(|c| PromptTokenUsageInfo { cached_tokens: c }),
337 }
338 }
339 +
340 + pub fn from_token_usage(usage: TokenUsage, enable_prompt_tokens_details: bool) -> Self {
341 + Self::from_counts(
342 + usage.prompt_token_count,
343 + usage.output_token_count,
344 + enable_prompt_tokens_details.then_some(usage.cached_token_count),
345 + )
346 + }
347 }
348
349 /// Mirrors the Python vLLM `PromptTokenUsageInfo` class.
335 -#[serde_with::skip_serializing_none]
350 #[derive(Debug, Clone, Serialize)]
351 pub struct PromptTokenUsageInfo {
338 - pub cached_tokens: Option<u32>,
352 + pub cached_tokens: usize,
353 +}
354 +
355 +#[cfg(test)]
356 +mod usage_tests {
357 + use vllm_llm::TokenUsage;
358 +
359 + use super::Usage;
360 +
361 + #[test]
362 + fn token_usage_hides_prompt_token_details_by_default() {
363 + let usage = Usage::from_token_usage(
364 + TokenUsage {
365 + prompt_token_count: 5,
366 + output_token_count: 2,
367 + cached_token_count: 3,
368 + },
369 + false,
370 + );
371 +
372 + assert_eq!(usage.prompt_tokens, 5);
373 + assert_eq!(usage.completion_tokens, Some(2));
374 + assert!(usage.prompt_tokens_details.is_none());
375 + }
376 +
377 + #[test]
378 + fn token_usage_includes_prompt_token_details_when_enabled() {
379 + let usage = Usage::from_token_usage(
380 + TokenUsage {
381 + prompt_token_count: 5,
382 + output_token_count: 2,
383 + cached_token_count: 3,
384 + },
385 + true,
386 + );
387 +
388 + assert_eq!(
389 + usage.prompt_tokens_details.as_ref().map(|details| details.cached_tokens),
390 + Some(3)
391 + );
392 + }
393 }
394
395 /// OpenAI completions-style logprobs.
rust/src/server/src/routes/tests.rs
+7 -2
@@ -43,6 +43,7 @@ use zeromq::prelude::{SocketRecv, SocketSend};
43 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
44
45 use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
46 +use crate::config::ApiServerOptions;
47 use crate::state::AppState;
48
49 fn request_output(
@@ -787,8 +788,12 @@ async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
788 )
789 .await;
790 let app = build_router(Arc::new(
790 - AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
791 - .with_request_id_headers(true),
791 + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_server_options(
792 + ApiServerOptions {
793 + enable_request_id_headers: true,
794 + ..Default::default()
795 + },
796 + ),
797 ));
798 (app, engine_task)
799 }
rust/src/server/src/routes/tokenize/types.rs
+3 -2
@@ -134,11 +134,12 @@ impl Normalizable for DetokenizeRequest {}
134
135 #[cfg(test)]
136 mod tests {
137 - use super::*;
138 - use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
137 use serde_json::json;
138 use vllm_chat::ChatTool;
139
140 + use super::*;
141 + use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
142 +
143 #[test]
144 fn tokenize_request_converts_openai_tools() {
145 // The untagged `TokenizeRequest` must resolve a messages+tools body to
rust/src/server/src/state.rs
+7 -15
@@ -9,6 +9,7 @@ use vllm_chat::ChatLlm;
9 use vllm_engine_core_client::EngineCoreClient;
10 use vllm_engine_core_client::protocol::lora::LoraRequest;
11
12 +use crate::config::ApiServerOptions;
13 use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
14 use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
15
@@ -27,10 +28,8 @@ pub struct AppState {
28 served_model_names: Vec<String>,
29 /// Shared chat facade used by all requests.
30 pub chat: ChatLlm,
30 - /// Whether to log a summary line for each completed request.
31 - pub enable_log_requests: bool,
32 - /// Whether to set X-Request-Id on every HTTP response.
33 - pub enable_request_id_headers: bool,
31 + /// HTTP/API-server behavior switches.
32 + pub api_server_options: ApiServerOptions,
33 /// Runtime server information returned by `/server_info`, when available.
34 server_info: Option<ServerInfoSnapshot>,
35 /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
@@ -58,8 +57,7 @@ impl AppState {
57 Self {
58 served_model_names,
59 chat,
61 - enable_log_requests: false,
62 - enable_request_id_headers: false,
60 + api_server_options: ApiServerOptions::default(),
61 server_info: None,
62 api_key_hashes: Vec::new(),
63 server_load: AtomicU64::new(0),
@@ -67,15 +65,9 @@ impl AppState {
65 }
66 }
67
70 - /// Enable per-request completion logging.
71 - pub fn with_log_requests(mut self, enabled: bool) -> Self {
72 - self.enable_log_requests = enabled;
73 - self
74 - }
75 -
76 - /// Enable X-Request-Id response headers.
77 - pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
78 - self.enable_request_id_headers = enabled;
68 + /// Set HTTP/API-server behavior switches.
69 + pub fn with_api_server_options(mut self, options: ApiServerOptions) -> Self {
70 + self.api_server_options = options;
71 self
72 }
73
rust/src/text/src/output/decoded.rs
+9 -5
@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
6 use tracing::{Level, debug, trace};
7 use vllm_engine_core_client::AbortCause;
8 use vllm_engine_core_client::protocol::StopReason;
9 -use vllm_llm::{FinishReason, GenerateOutput};
9 +use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
10 use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
11
12 use super::logprobs::{
@@ -40,8 +40,7 @@ impl Default for TextDecodeOptions {
40 /// Terminal metadata carried on the final [`DecodedTextEvent`].
41 #[derive(Debug, Clone, PartialEq)]
42 pub struct Finished {
43 - pub prompt_token_count: usize,
44 - pub output_token_count: usize,
43 + pub usage: TokenUsage,
44 pub finish_reason: FinishReason,
45 /// Connector-specific KV transfer parameters for disaggregated serving.
46 pub kv_transfer_params: Option<serde_json::Value>,
@@ -98,12 +97,14 @@ pub async fn decoded_text_event_stream(
97 ) -> crate::Result<()> {
98 let mut decoder: Option<Box<dyn IncrementalDecoder>> = None;
99 let mut prompt_token_count = 0_usize;
100 + let mut cached_token_count = 0_usize;
101 let mut token_ids = Vec::new();
102 let mut output_token_count: usize = 0;
103 let mut logprobs: Option<DecodedLogprobs> = None;
104
105 while let Some(next) = raw_stream.next().await {
106 let output = next?;
107 + cached_token_count = cached_token_count.max(output.cached_token_count);
108
109 // If it's the first output, init states and yield `Start` event.
110 if decoder.is_none() {
@@ -267,8 +268,11 @@ pub async fn decoded_text_event_stream(
268 token_ids,
269 logprobs,
270 finished: Some(Finished {
270 - prompt_token_count,
271 - output_token_count,
271 + usage: TokenUsage {
272 + prompt_token_count,
273 + output_token_count,
274 + cached_token_count,
275 + },
276 finish_reason: reason,
277 kv_transfer_params,
278 }),
rust/src/text/src/output/mod.rs
+13 -4
@@ -23,6 +23,7 @@ pub struct CollectedTextOutput {
23 pub logprobs: Option<DecodedLogprobs>,
24 pub token_ids: Vec<u32>,
25 pub finish_reason: FinishReason,
26 + pub usage: vllm_llm::TokenUsage,
27 /// Connector-specific KV transfer parameters for disaggregated serving.
28 pub kv_transfer_params: Option<serde_json::Value>,
29 }
@@ -74,6 +75,7 @@ impl<T: TextOutputStream> T {
75 logprobs: delta_logprobs,
76 token_ids: delta_token_ids,
77 finish_reason: FinishReason::Error,
78 + usage: vllm_llm::TokenUsage::default(),
79 kv_transfer_params: None,
80 })
81 };
@@ -81,6 +83,7 @@ impl<T: TextOutputStream> T {
83 if let Some(finished) = finished {
84 let mut collected = collected.unwrap();
85 collected.finish_reason = finished.finish_reason;
86 + collected.usage = finished.usage;
87 collected.kv_transfer_params = finished.kv_transfer_params;
88 return Ok(collected);
89 }
@@ -146,8 +149,11 @@ mod tests {
149 ],
150 }),
151 finished: Some(Finished {
149 - prompt_token_count: 2,
150 - output_token_count: 2,
152 + usage: vllm_llm::TokenUsage {
153 + prompt_token_count: 2,
154 + output_token_count: 2,
155 + cached_token_count: 0,
156 + },
157 finish_reason: FinishReason::stop_eos(),
158 kv_transfer_params: None,
159 }),
@@ -260,8 +266,11 @@ mod tests {
266 ],
267 }),
268 finished: Some(Finished {
263 - prompt_token_count: 2,
264 - output_token_count: 5,
269 + usage: vllm_llm::TokenUsage {
270 + prompt_token_count: 2,
271 + output_token_count: 5,
272 + cached_token_count: 0,
273 + },
274 finish_reason: FinishReason::stop_eos(),
275 kv_transfer_params: None,
276 }),