[Bugfix][Rust Frontend] Use server-reported prompt_tokens in vllm-bench (#59694)

rust/src/bench/src/backends/mod.rs

@@ -52,6 +52,21 @@ pub struct ChatDelta {
 #[derive(Deserialize)]
 pub struct ChunkUsage {
     pub completion_tokens: Option<u64>,
+    pub prompt_tokens: Option<u64>,
+}
+
+impl ChunkUsage {
+    /// Record the server-reported token counts on `output`. A zero
+    /// `prompt_tokens` keeps the client estimate: some OpenAI-compatible
+    /// servers send placeholder zeros in usage chunks.
+    pub fn apply(&self, output: &mut RequestFuncOutput) {
+        if let Some(ct) = self.completion_tokens {
+            output.output_tokens = ct as usize;
+        }
+        if let Some(pt) = self.prompt_tokens.filter(|&pt| pt > 0) {
+            output.prompt_len = pt as usize;
+        }
+    }
 }
 
 use crate::cli::BackendKind;
@@ -246,6 +261,21 @@ pub fn build_headers(
 mod tests {
     use super::*;
 
+    /// A placeholder zero `prompt_tokens` must not replace the client estimate.
+    #[test]
+    fn chunk_usage_ignores_zero_prompt_tokens() {
+        let mut output = RequestFuncOutput {
+            prompt_len: 100,
+            ..Default::default()
+        };
+        ChunkUsage {
+            completion_tokens: Some(3),
+            prompt_tokens: Some(0),
+        }
+        .apply(&mut output);
+        assert_eq!((output.prompt_len, output.output_tokens), (100, 3));
+    }
+
     #[test]
     fn test_update_server_metrics_converts_ms_to_seconds() {
         let chunk: CompletionChunk = serde_json::from_value(serde_json::json!({

rust/src/bench/src/backends/openai_chat.rs

@@ -130,10 +130,8 @@ impl OpenAIChatBackend {
                             }
                             // Separate `if` (not `else if`) — Dynamo may send
                             // both choices and usage in the same chunk.
-                            if let Some(ref usage) = data.usage
-                                && let Some(ct) = usage.completion_tokens
-                            {
-                                output.output_tokens = ct as usize;
+                            if let Some(usage) = &data.usage {
+                                usage.apply(&mut output);
                             }
                         }
                     }
@@ -363,7 +361,11 @@ mod tests {
     /// Serve one streaming chat response over a raw socket: `num_tokens`
     /// content chunks back-to-back, then `usage_delay` of silence, then the
     /// choice-less usage chunk and `[DONE]`. Returns the endpoint URL.
-    fn spawn_sse_server(num_tokens: usize, usage_delay: std::time::Duration) -> String {
+    fn spawn_sse_server(
+        num_tokens: usize,
+        usage_delay: std::time::Duration,
+        prompt_tokens: Option<u64>,
+    ) -> String {
         use std::io::{BufRead, BufReader, Read, Write};
 
         let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
@@ -393,8 +395,10 @@ mod tests {
                 send("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n");
             }
             std::thread::sleep(usage_delay);
+            let prompt_tokens =
+                prompt_tokens.map_or(String::new(), |pt| format!(",\"prompt_tokens\":{pt}"));
             send(&format!(
-                "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}}}}}\n\n"
+                "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}{prompt_tokens}}}}}\n\n"
             ));
             send("data: [DONE]\n\n");
         });
@@ -411,7 +415,7 @@ mod tests {
         // real delay is needed here; there is no observable event to wait on.
         let usage_delay = std::time::Duration::from_millis(150);
         let input = RequestFuncInput {
-            api_url: spawn_sse_server(3, usage_delay),
+            api_url: spawn_sse_server(3, usage_delay, None),
             model: "test-model".to_string(),
             output_len: 3,
             ..Default::default()
@@ -442,4 +446,42 @@ mod tests {
         // keys already set above must not be overridden by extra_body
         assert_eq!(v["stream"], true);
     }
+
+    /// The server's `usage.prompt_tokens` replaces the client-side estimate,
+    /// so input-token totals count what the server actually processed
+    /// (chat template, multimodal tokens). Mirrors Python #38654.
+    #[tokio::test]
+    async fn test_usage_prompt_tokens_override_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(3, std::time::Duration::ZERO, Some(128)),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 3,
+            ..Default::default()
+        };
+
+        let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!(output.prompt_len, 128);
+    }
+
+    /// Servers that omit `usage.prompt_tokens` keep the client-side estimate,
+    /// and `completion_tokens` is still recorded.
+    #[tokio::test]
+    async fn test_missing_prompt_tokens_keeps_client_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(3, std::time::Duration::ZERO, None),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 3,
+            ..Default::default()
+        };
+
+        let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!(output.prompt_len, 100);
+        assert_eq!(output.output_tokens, 3);
+    }
 }

rust/src/bench/src/backends/openai_completions.rs

@@ -156,10 +156,8 @@ impl OpenAICompletionsBackend {
                             }
                             // Separate `if` (not `else if`) — Dynamo may send
                             // both choices and usage in the same chunk.
-                            if let Some(ref usage) = data.usage
-                                && let Some(ct) = usage.completion_tokens
-                            {
-                                output.output_tokens = ct as usize;
+                            if let Some(usage) = &data.usage {
+                                usage.apply(&mut output);
                             }
                         }
                     }
@@ -201,3 +199,69 @@ impl OpenAICompletionsBackend {
         Ok(output)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Serve `body` as one streaming completions response over a raw socket.
+    /// Returns the endpoint URL.
+    fn spawn_sse_server(body: &'static str) -> String {
+        use std::io::{BufRead, BufReader, Read, Write};
+
+        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let addr = listener.local_addr().unwrap();
+        std::thread::spawn(move || {
+            let (mut stream, _) = listener.accept().unwrap();
+            let mut reader = BufReader::new(stream.try_clone().unwrap());
+            let mut content_length = 0;
+            loop {
+                let mut line = String::new();
+                reader.read_line(&mut line).unwrap();
+                if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
+                    content_length = v.trim().parse().unwrap();
+                }
+                if line == "\r\n" {
+                    break;
+                }
+            }
+            reader.read_exact(&mut vec![0; content_length]).unwrap();
+            stream
+                .write_all(
+                    format!(
+                        "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\
+                         connection: close\r\n\r\n{body}"
+                    )
+                    .as_bytes(),
+                )
+                .unwrap();
+        });
+        format!("http://{addr}/v1/completions")
+    }
+
+    /// The server's `usage.prompt_tokens` replaces the client-side estimate,
+    /// as in the chat backend.
+    #[tokio::test]
+    async fn test_usage_prompt_tokens_override_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(concat!(
+                "data: {\"choices\":[{\"text\":\"a\"}]}\n\n",
+                "data: {\"choices\":[{\"text\":\"b\"}]}\n\n",
+                "data: {\"choices\":[],\"usage\":{\"completion_tokens\":2,\"prompt_tokens\":128}}\n\n",
+                "data: [DONE]\n\n",
+            )),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 2,
+            ..Default::default()
+        };
+
+        let output = OpenAICompletionsBackend
+            .send_request(&input, &reqwest::Client::new())
+            .await
+            .unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!((output.prompt_len, output.output_tokens), (128, 2));
+    }
+}

rust/src/bench/src/backends/openai_responses.rs

@@ -27,6 +27,7 @@ struct ResponsesEventResponse {
 
 #[derive(serde::Deserialize)]
 struct ResponsesUsage {
+    input_tokens: Option<u64>,
     output_tokens: Option<u64>,
 }
 
@@ -180,8 +181,14 @@ impl ResponsesStream {
             // with `status: "incomplete"`; the spec also has `response.incomplete`.
             Some("response.completed" | "response.incomplete") => {
                 self.saw_terminal_event = true;
-                if let Some(n) = event.response.and_then(|r| r.usage?.output_tokens) {
-                    self.output.output_tokens = n as usize;
+                if let Some(usage) = event.response.and_then(|r| r.usage) {
+                    if let Some(n) = usage.output_tokens {
+                        self.output.output_tokens = n as usize;
+                    }
+                    // Same rule as `ChunkUsage::apply`: a zero keeps the client estimate.
+                    if let Some(n) = usage.input_tokens.filter(|&n| n > 0) {
+                        self.output.prompt_len = n as usize;
+                    }
                 }
             }
             kind if matches!(kind, Some("error" | "response.failed")) || event.error.is_some() => {
@@ -253,6 +260,21 @@ mod tests {
         stream.finish()
     }
 
+    /// The terminal event's `usage.input_tokens` replaces the client estimate,
+    /// as in Python's `async_request_openai_responses`.
+    #[test]
+    fn test_terminal_usage_input_tokens_override_prompt_len() {
+        let output = run(&[
+            text_delta("Hi"),
+            event(
+                "response.completed",
+                ",\"response\":{\"usage\":{\"input_tokens\":128,\"output_tokens\":1}}",
+            ),
+        ]);
+        assert!(output.success, "{}", output.error);
+        assert_eq!((output.prompt_len, output.output_tokens), (128, 1));
+    }
+
     #[test]
     fn test_reasoning_starts_ttft_and_latency_ends_at_last_token() {
         let output = run(&[

rust/src/bench/src/benchmark.rs

@@ -645,7 +645,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum();
         println!("Dry run stats:");
         println!("  Total prompts: {}", input_requests.len());
-        println!("  Total input tokens: {total_input_tokens}");
+        println!("  Estimated total input tokens: {total_input_tokens}");
         println!("  Total expected output tokens: {total_output_tokens}");
         println!(
             "  Avg input tokens: {:.1}",
@@ -1118,7 +1118,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         (m, Vec::new())
     } else {
         calculate_metrics(
-            &input_requests,
             &outputs,
             benchmark_duration,
             &config.selected_percentiles,
@@ -1148,7 +1147,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         ) {
             let ss = steady_state::compute(
                 &outputs,
-                &input_requests,
                 &window,
                 &config.selected_percentiles,
                 config.backend.is_pooling(),

rust/src/bench/src/metrics/calculator.rs

@@ -3,7 +3,6 @@
 
 use crate::backends::RequestFuncOutput;
 use crate::config::GoodputConfig;
-use crate::datasets::SampleRequest;
 use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
 use crate::multi_turn::ConversationOutput;
 
@@ -27,7 +26,6 @@ fn log_failed_requests(outputs: &[RequestFuncOutput]) {
 ///
 /// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
 pub fn calculate_metrics(
-    input_requests: &[SampleRequest],
     outputs: &[RequestFuncOutput],
     dur_s: f64,
     selected_percentiles: &[f64],
@@ -44,7 +42,7 @@ pub fn calculate_metrics(
     let mut ttfts: Vec<f64> = Vec::new();
     let mut e2els: Vec<f64> = Vec::new();
 
-    for (i, output) in outputs.iter().enumerate() {
+    for output in outputs.iter() {
         if output.success {
             let output_len = if output.output_tokens > 0 {
                 output.output_tokens
@@ -57,7 +55,7 @@ pub fn calculate_metrics(
             };
 
             actual_output_lens.push(output_len);
-            total_input += input_requests[i].prompt_len;
+            total_input += output.prompt_len;
 
             if output_len > 1 {
                 let latency_minus_ttft = output.latency - output.ttft;
@@ -283,7 +281,6 @@ fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64
 ///
 /// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py.
 /// Key differences from `calculate_metrics()`:
-/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len`
 /// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming)
 /// - `total_output` is 0 (pooling produces no output tokens)
 /// - `total_token_throughput` = `total_input / dur_s` (input tokens only)
@@ -380,7 +377,7 @@ pub fn calculate_embedding_metrics(
 
 /// Calculate multi-turn benchmark metrics with per-turn breakdown.
 ///
-/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics
+/// 1. Flatten all turns' outputs → overall metrics
 /// 2. Group by turn_index → per-turn metrics
 /// 3. Conversation-level stats
 pub fn calculate_multi_turn_metrics(
@@ -391,24 +388,15 @@ pub fn calculate_multi_turn_metrics(
     max_turn_count: usize,
 ) -> MultiTurnMetrics {
     // Flatten all turns for overall metrics
-    let mut all_requests: Vec<SampleRequest> = Vec::new();
     let mut all_outputs: Vec<RequestFuncOutput> = Vec::new();
 
     for conv in conversation_outputs {
         for turn in &conv.turns {
-            all_requests.push(SampleRequest {
-                prompt: std::sync::Arc::from(""),
-                prompt_len: turn.request_output.prompt_len,
-                expected_output_len: turn.request_output.output_tokens,
-                request_id: None,
-                ..Default::default()
-            });
             all_outputs.push(turn.request_output.clone());
         }
     }
 
     let (overall, _) = calculate_metrics(
-        &all_requests,
         &all_outputs,
         dur_s,
         selected_percentiles,
@@ -419,25 +407,16 @@ pub fn calculate_multi_turn_metrics(
     // Per-turn breakdown
     let mut per_turn: Vec<BenchmarkMetrics> = Vec::new();
     for turn_idx in 0..max_turn_count {
-        let mut turn_requests: Vec<SampleRequest> = Vec::new();
         let mut turn_outputs: Vec<RequestFuncOutput> = Vec::new();
 
         for conv in conversation_outputs {
             if let Some(turn) = conv.turns.get(turn_idx) {
-                turn_requests.push(SampleRequest {
-                    prompt: std::sync::Arc::from(""),
-                    prompt_len: turn.request_output.prompt_len,
-                    expected_output_len: turn.request_output.output_tokens,
-                    request_id: None,
-                    ..Default::default()
-                });
                 turn_outputs.push(turn.request_output.clone());
             }
         }
 
         if !turn_outputs.is_empty() {
             let (turn_metrics, _) = calculate_metrics(
-                &turn_requests,
                 &turn_outputs,
                 dur_s,
                 selected_percentiles,
@@ -509,6 +488,26 @@ mod tests {
         assert!((p99 - 98.01).abs() < 0.1);
     }
 
+    /// Input-token totals come from each output's `prompt_len` (server-reported
+    /// via `usage.prompt_tokens`), not the client-side request estimate.
+    /// Mirrors Python's `serve.py`.
+    #[test]
+    fn test_total_input_uses_output_prompt_len() {
+        let outputs = vec![
+            RequestFuncOutput {
+                success: true,
+                prompt_len: 128,
+                output_tokens: 3,
+                ..Default::default()
+            };
+            2
+        ];
+
+        let (metrics, _) = calculate_metrics(&outputs, 1.0, &[], false, &GoodputConfig::default());
+
+        assert_eq!(metrics.total_input, 256);
+    }
+
     #[test]
     fn test_multi_turn_metrics_use_request_prompt_lengths() {
         let turn = |turn_index, prompt_len, output_tokens| crate::multi_turn::TurnOutput {

rust/src/bench/src/metrics/steady_state.rs

@@ -170,16 +170,12 @@ pub fn detect_window(
     })
 }
 
-use crate::datasets::SampleRequest;
-
-/// Compute steady-state metrics from per-request outputs, the detected window,
-/// and the original request list (for input-token attribution).
+/// Compute steady-state metrics from per-request outputs and the detected window.
 ///
 /// `is_pooling` gates off TTFT because pooling backends write `ttft = latency`
 /// as a placeholder; reporting TTFT over that would silently be E2EL.
 pub fn compute(
     outputs: &[RequestFuncOutput],
-    input_requests: &[SampleRequest],
     window: &SteadyStateWindow,
     percentiles: &[f64],
     is_pooling: bool,
@@ -200,7 +196,7 @@ pub fn compute(
     let mut ttfts_in_window: Vec<f64> = Vec::new();
     let mut tpots_in_window: Vec<f64> = Vec::new();
 
-    for (idx, o) in outputs.iter().enumerate() {
+    for o in outputs.iter() {
         if !o.success {
             continue;
         }
@@ -213,7 +209,7 @@ pub fn compute(
         }
 
         if started_in {
-            input_tokens_in_window += input_requests[idx].prompt_len;
+            input_tokens_in_window += o.prompt_len;
             if !is_pooling {
                 ttfts_in_window.push(o.ttft);
                 // Per-request TPOT: (latency - ttft) / (output_tokens - 1).
@@ -417,10 +413,6 @@ mod tests {
         assert!((w.end_s - 10.0).abs() < 1e-9);
     }
 
-    use std::sync::Arc;
-
-    use crate::datasets::SampleRequest;
-
     fn mk_full(
         start: f64,
         latency: f64,
@@ -440,16 +432,6 @@ mod tests {
         }
     }
 
-    fn mk_sample(prompt_len: usize) -> SampleRequest {
-        SampleRequest {
-            prompt: Arc::from(""),
-            prompt_len,
-            expected_output_len: 0,
-            request_id: None,
-            ..Default::default()
-        }
-    }
-
     #[test]
     fn compute_constant_concurrency_matches_full_run() {
         // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s.
@@ -459,13 +441,10 @@ mod tests {
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
 
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
         let m = compute(
             &outputs,
-            &requests,
             &w,
             &[99.0],
             false, // is_pooling
@@ -502,12 +481,10 @@ mod tests {
         let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s
         let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000);
         outputs.push(late);
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
 
         // Target = 2 so threshold_abs = 2; window opens when both are in flight.
         let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
 
         // Late request started in window but most of its tokens are emitted AFTER end_s.
         // Only a handful (≤3) should count.
@@ -534,10 +511,8 @@ mod tests {
             mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
             mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true);
+        let m = compute(&outputs, &w, &[99.0], /* is_pooling */ true);
         assert_eq!(m.mean_ttft_ms, 0.0);
         assert_eq!(m.median_ttft_ms, 0.0);
         assert!(m.percentiles_ttft_ms.is_empty());
@@ -560,10 +535,8 @@ mod tests {
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
 
         assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6);
         assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6);
@@ -574,10 +547,8 @@ mod tests {
     #[test]
     fn json_roundtrip_some() {
         let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
         let s = serde_json::to_string(&m).unwrap();
         let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap();
         assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12);
@@ -605,14 +576,13 @@ mod tests {
         let itl = vec![1.0; 9];
         let outputs: Vec<RequestFuncOutput> =
             (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect();
-        let requests: Vec<SampleRequest> = (0..5).map(|_| mk_sample(100)).collect();
 
         let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap();
         assert_eq!(w.threshold_abs, 5);
         assert!((w.start_s - 4.0).abs() < 1e-9);
         assert!((w.end_s - 20.0).abs() < 1e-9);
 
-        let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false);
+        let m = compute(&outputs, &w, &[50.0, 99.0], false);
 
         // Round-trip through JSON.
         let s = serde_json::to_string(&m).unwrap();