Statistics for reasoning and prefix-cached tokens (#413)

Co-authored-by: Venkat Raman <vraman2811@gmail.com>

Author: guoqingbao

docs/prefix_cache.md

@@ -17,6 +17,30 @@ Prefix cache is block-granular. If a shared prefix ends in the middle of a block
 
 If `--prefix-cache-max-tokens` is omitted, the cache defaults to roughly 25% of GPU KV blocks in this project.
 
+## Usage Reporting
+
+OpenAI-compatible chat responses include prefix-cache and reasoning token details when they are non-zero:
+
+```json
+{
+  "usage": {
+    "prompt_tokens": 128,
+    "completion_tokens": 64,
+    "total_tokens": 192,
+    "prompt_time_costs": 4,
+    "completion_time_costs": 250,
+    "prompt_tokens_details": {
+      "cached_tokens": 64
+    },
+    "completion_tokens_details": {
+      "reasoning_tokens": 32
+    }
+  }
+}
+```
+
+`prompt_tokens_details.cached_tokens` reports the number of prompt tokens reused from the prefix cache. `completion_tokens_details.reasoning_tokens` reports generated tokens inside reasoning blocks such as `<think>...</think>`. Both detail objects are omitted when their count is zero.
+
 ## Hybrid Mamba snapshot stride
 
 For hybrid Mamba models, prefix reuse also needs compatible snapshot boundaries.

src/openai/pipelines/llm_engine.rs

@@ -16,7 +16,8 @@ use crate::openai::streaming::ChatResponse;
 use crate::openai::TaskData;
 use crate::scheduler::Scheduler;
 use crate::tools::stream_parser::{
-    BufferedFinalizeResult, ParserState, StreamResult, StreamToolParser,
+    extract_reasoning_content, strip_reasoning_markers, BufferedFinalizeResult, ParserState,
+    StreamResult, StreamToolParser,
 };
 #[cfg(feature = "flashinfer")]
 use crate::FlashInferKvParams;
@@ -27,8 +28,8 @@ use crate::{
         multimodal::ImageData,
         responses::{
             ChatChoice, ChatChoiceData, ChatCompletionChunk, ChatCompletionUsageResponse, Choice,
-            ChoiceData, EmbeddingData, EmbeddingOutput, EmbeddingResponse, EmbeddingUsage,
-            WrapperLogprobs,
+            ChoiceData, CompletionTokensDetails, EmbeddingData, EmbeddingOutput, EmbeddingResponse,
+            EmbeddingUsage, PromptTokensDetails, WrapperLogprobs,
         },
         sampling_params::Logprobs,
         sampling_params::SamplingParams,
@@ -860,6 +861,16 @@ impl LLMEngine {
                     }
 
                     //chat completion statistics
+                    let cached_tokens_total = result
+                        .values()
+                        .filter_map(|(_, usage)| usage.prompt_tokens_details.as_ref())
+                        .map(|details| details.cached_tokens)
+                        .sum();
+                    let reasoning_tokens_total = result
+                        .values()
+                        .filter_map(|(_, usage)| usage.completion_tokens_details.as_ref())
+                        .map(|details| details.reasoning_tokens)
+                        .sum();
                     let overall_usage = ChatCompletionUsageResponse {
                         request_id: "".to_string(),
                         created: 0,
@@ -878,6 +889,15 @@ impl LLMEngine {
                             .map(|(_, usage)| usage.completion_time_costs)
                             .max()
                             .unwrap_or(0),
+                        prompt_tokens_details: (cached_tokens_total > 0)
+                            .then_some(PromptTokensDetails {
+                                cached_tokens: cached_tokens_total,
+                            }),
+                        completion_tokens_details: (reasoning_tokens_total > 0).then_some(
+                            CompletionTokensDetails {
+                                reasoning_tokens: reasoning_tokens_total,
+                            },
+                        ),
                     };
 
                     let prompt_tps : f32 = result.values().map(|(_, usage)| {
@@ -1270,6 +1290,66 @@ impl LLMEngine {
         self.scheduler.has_unfinished_sequences()
     }
 
+    fn count_text_tokens(pipeline: &DefaultPipeline, text: &str) -> usize {
+        if text.is_empty() {
+            return 0;
+        }
+        pipeline
+            .tokenizer()
+            .encode(text, false)
+            .map(|encoding| encoding.get_ids().len())
+            .unwrap_or(0)
+    }
+
+    fn reasoning_token_count_for_sequence(
+        &self,
+        pipeline: &DefaultPipeline,
+        group: &SequenceGroup,
+        seq: &Arc<Sequence>,
+    ) -> usize {
+        let outputs = seq.deref().get_output_tokens();
+        if outputs.is_empty() {
+            return 0;
+        }
+
+        let raw_output = outputs.iter().map(|x| x.bytes.as_str()).collect::<String>();
+        if let Some((reasoning, _)) = extract_reasoning_content(&raw_output) {
+            return Self::count_text_tokens(pipeline, &reasoning);
+        }
+
+        let should_parse_tools = group.sampling_params.mcp_mode.is_some();
+        if !(crate::stream_as_reasoning_content() && should_parse_tools) {
+            return 0;
+        }
+
+        let mut parser = StreamToolParser::new_with_config(
+            &pipeline.tool_model_type,
+            pipeline.tool_parser_model_id.clone(),
+            pipeline.tool_config.clone(),
+            group.tools.clone(),
+            pipeline.enforce_parser.clone(),
+        );
+        if group.active_reasoning_end.is_some() {
+            parser.set_initial_reasoning_end_marker(group.active_reasoning_end.clone());
+        }
+        parser.set_detect_tools_in_reasoning(true);
+
+        let mut reasoning_text = String::new();
+        for output in outputs {
+            let was_in_reasoning = parser.in_reasoning();
+            match parser.process_token(output.token, &output.bytes) {
+                StreamResult::Content(text) | StreamResult::FlushBuffer(text) => {
+                    if was_in_reasoning {
+                        reasoning_text.push_str(&strip_reasoning_markers(&text));
+                    }
+                }
+                StreamResult::Buffering | StreamResult::ToolCalls(_) => {}
+            }
+        }
+
+        Self::count_text_tokens(pipeline, &reasoning_text)
+    }
+
     fn schedule_current_batch(&mut self, rank: usize) -> Result<()> {
         let scheduler_outputs = self.scheduler.schedule();
         if !scheduler_outputs.ignored_seq_groups.is_empty() {
@@ -1704,9 +1784,9 @@ impl LLMEngine {
 
                 let mut choices = Vec::new();
                 let do_sync_response = allow_sync_response && group.sender.is_none();
+                let pipeline = self.get_pipeline(0usize).unwrap().0.as_ref();
 
                 if do_sync_response {
-                    let pipeline = self.get_pipeline(0usize).unwrap().0.as_ref();
                     for (index, seq) in top_n.iter().enumerate() {
                         let outputs = seq.deref_mut().get_output_tokens();
                         let should_parse_tools = group.sampling_params.mcp_mode.is_some();
@@ -1929,6 +2009,14 @@ impl LLMEngine {
                     .duration_since(group.created_time)
                     .unwrap()
                     .as_millis();
+                let cached_tokens = top_n
+                    .first()
+                    .and_then(|seq| self.get_num_cached_tokens_for_seq(seq.deref().get_id()))
+                    .unwrap_or_else(|| top_n.first().unwrap().deref().get_num_cached_tokens());
+                let reasoning_tokens = top_n
+                    .iter()
+                    .map(|seq| self.reasoning_token_count_for_sequence(pipeline, group, seq))
+                    .sum::<usize>();
 
                 let usage = ChatCompletionUsageResponse {
                     request_id: group.request_id.clone(),
@@ -1938,6 +2026,10 @@ impl LLMEngine {
                     total_tokens: completion_tokens + prompt_tokens,
                     prompt_time_costs: prompt_time_costs as usize,
                     completion_time_costs: completion_time_costs as usize,
+                    prompt_tokens_details: (cached_tokens > 0)
+                        .then_some(PromptTokensDetails { cached_tokens }),
+                    completion_tokens_details: (reasoning_tokens > 0)
+                        .then_some(CompletionTokensDetails { reasoning_tokens }),
                 };
 
                 responses.insert(group.request_id.clone(), (choices, usage.clone()));
@@ -2105,4 +2197,8 @@ impl LLMEngine {
     pub fn query_prefix_cache_match_tokens(&mut self, tokens: &[u32]) -> usize {
         self.scheduler.query_prefix_cache_match_tokens(tokens)
     }
+
+    pub fn get_num_cached_tokens_for_seq(&self, seq_id: usize) -> Option<usize> {
+        self.scheduler.get_num_cached_tokens_for_seq(seq_id)
+    }
 }

src/openai/responses.rs

@@ -39,6 +39,16 @@ macro_rules! try_api {
     };
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PromptTokensDetails {
+    pub cached_tokens: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CompletionTokensDetails {
+    pub reasoning_tokens: usize,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ChatCompletionUsageResponse {
     pub request_id: String,
@@ -48,6 +58,10 @@ pub struct ChatCompletionUsageResponse {
     pub total_tokens: usize,
     pub prompt_time_costs: usize,     //milliseconds
     pub completion_time_costs: usize, //milliseconds
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub prompt_tokens_details: Option<PromptTokensDetails>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub completion_tokens_details: Option<CompletionTokensDetails>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -188,3 +202,66 @@ pub struct EmbeddingResponse {
     pub model: String,
     pub usage: EmbeddingUsage,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{ChatCompletionUsageResponse, CompletionTokensDetails, PromptTokensDetails};
+
+    fn usage(
+        prompt_tokens_details: Option<PromptTokensDetails>,
+        completion_tokens_details: Option<CompletionTokensDetails>,
+    ) -> ChatCompletionUsageResponse {
+        ChatCompletionUsageResponse {
+            request_id: "request-id".to_string(),
+            created: 0,
+            completion_tokens: 50,
+            prompt_tokens: 100,
+            total_tokens: 150,
+            prompt_time_costs: 1,
+            completion_time_costs: 1,
+            prompt_tokens_details,
+            completion_tokens_details,
+        }
+    }
+
+    #[test]
+    fn usage_omits_token_details_when_none() {
+        let value = serde_json::to_value(usage(None, None)).expect("serialize usage");
+        let object = value.as_object().expect("usage is a JSON object");
+
+        assert!(!object.contains_key("prompt_tokens_details"));
+        assert!(!object.contains_key("completion_tokens_details"));
+    }
+
+    #[test]
+    fn usage_includes_prompt_tokens_details_when_some() {
+        let value =
+            serde_json::to_value(usage(Some(PromptTokensDetails { cached_tokens: 64 }), None))
+                .expect("serialize usage");
+
+        assert_eq!(
+            value
+                .pointer("/prompt_tokens_details/cached_tokens")
+                .and_then(|v| v.as_u64()),
+            Some(64)
+        );
+    }
+
+    #[test]
+    fn usage_includes_completion_tokens_details_when_some() {
+        let value = serde_json::to_value(usage(
+            None,
+            Some(CompletionTokensDetails {
+                reasoning_tokens: 32,
+            }),
+        ))
+        .expect("serialize usage");
+
+        assert_eq!(
+            value
+                .pointer("/completion_tokens_details/reasoning_tokens")
+                .and_then(|v| v.as_u64()),
+            Some(32)
+        );
+    }
+}

src/scheduler/mod.rs

@@ -34,6 +34,7 @@ use self::{
 };
 
 const PREFIX_CACHE_PRESSURE_EVICT_PERCENT: f32 = 0.1; // evict 10% of prefix cache when under pressure
+const FINISHED_CACHED_TOKENS_MAX: usize = 16_384;
 
 pub struct SchedulerOutput {
     pub scheduled: Arc<VecDeque<Arc<SequenceGroup>>>,
@@ -57,6 +58,7 @@ pub struct Scheduler {
     mamba_state: MambaState,
     is_last_prefill: bool,
     prefill_chunk_size: usize,
+    finished_cached_tokens: HashMap<usize, usize>,
 }
 
 impl Scheduler {
@@ -84,6 +86,7 @@ impl Scheduler {
             mamba_state: MambaState::default(),
             is_last_prefill: false,
             prefill_chunk_size,
+            finished_cached_tokens: HashMap::new(),
         }
     }
 
@@ -282,6 +285,7 @@ impl Scheduler {
         for group in to_free {
             for seq in group.get_seqs().values() {
                 let seq_id = seq.deref().get_id();
+                self.remember_finished_cached_tokens(seq_id, seq.deref().get_num_cached_tokens());
                 let full_blocks = seq.deref().get_len() / self.block_engine.get_block_size();
                 let block_id = self
                     .block_engine
@@ -297,6 +301,32 @@ impl Scheduler {
         released_ids
     }
 
+    pub fn get_num_cached_tokens_for_seq(&self, seq_id: usize) -> Option<usize> {
+        self.running
+            .iter()
+            .chain(self.waiting.iter())
+            .chain(self.swapped_out.iter())
+            .find_map(|group| {
+                group
+                    .get_seqs()
+                    .values()
+                    .find(|seq| seq.deref().get_id() == seq_id)
+                    .map(|seq| seq.deref().get_num_cached_tokens())
+            })
+            .or_else(|| self.finished_cached_tokens.get(&seq_id).copied())
+    }
+
+    fn remember_finished_cached_tokens(&mut self, seq_id: usize, num_cached_tokens: usize) {
+        self.finished_cached_tokens
+            .insert(seq_id, num_cached_tokens);
+        while self.finished_cached_tokens.len() > FINISHED_CACHED_TOKENS_MAX {
+            let Some(oldest_seq_id) = self.finished_cached_tokens.keys().min().copied() else {
+                break;
+            };
+            self.finished_cached_tokens.remove(&oldest_seq_id);
+        }
+    }
+
     pub fn prefix_cache_enabled(&self) -> bool {
         self.block_engine.prefix_cache_enabled()
     }