[Bugfix][Rust Frontend] Stop vllm-bench chat latency at the last token (#59251)
Signed-off-by: zhenchenglin <zhl132@ucsd.edu> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Roy Wang <jasonailu87@gmail.com>
ZhenchengLin committed
Sep 30, 2026 at 19:04 UTC
31ec1c3aa18a4881889cf4bacd50a867a8248717
1 file changed
+73
-6
rust/src/bench/src/backends/openai_chat.rs
+73
-6
@@ -96,10 +96,6 @@ impl OpenAIChatBackend {
96
continue;
97
}
98
99
- // Python chat backend: timestamp is captured for ALL
100
- // non-DONE messages, and most_recent_timestamp is updated
101
- // unconditionally (outside `if choices:`). This differs from
102
- // completions which only timestamps content chunks.
99
let timestamp = Instant::now();
100
101
let data: ChatChunk = match serde_json::from_str(chunk) {
@@ -126,6 +122,9 @@ impl OpenAIChatBackend {
122
}
123
124
generated_text.push_str(content);
125
+ // Only token chunks advance the request end;
126
+ // the trailing usage chunk carries no token.
127
+ most_recent_timestamp = timestamp;
128
}
129
// Separate `if` (not `else if`) — Dynamo may send
130
// both choices and usage in the same chunk.
@@ -134,8 +133,6 @@ impl OpenAIChatBackend {
133
{
134
output.output_tokens = ct as usize;
135
}
137
-
138
- most_recent_timestamp = timestamp;
136
}
137
}
138
@@ -361,6 +358,76 @@ mod tests {
358
assert_eq!(a, b);
359
}
360
361
+ /// Serve one streaming chat response over a raw socket: `num_tokens`
362
+ /// content chunks back-to-back, then `usage_delay` of silence, then the
363
+ /// choice-less usage chunk and `[DONE]`. Returns the endpoint URL.
364
+ fn spawn_sse_server(num_tokens: usize, usage_delay: std::time::Duration) -> String {
365
+ use std::io::{BufRead, BufReader, Read, Write};
366
+
367
+ let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
368
+ let addr = listener.local_addr().unwrap();
369
+ std::thread::spawn(move || {
370
+ let (mut stream, _) = listener.accept().unwrap();
371
+ let mut reader = BufReader::new(stream.try_clone().unwrap());
372
+ let mut content_length = 0;
373
+ loop {
374
+ let mut line = String::new();
375
+ reader.read_line(&mut line).unwrap();
376
+ if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
377
+ content_length = v.trim().parse().unwrap();
378
+ }
379
+ if line == "\r\n" {
380
+ break;
381
+ }
382
+ }
383
+ reader.read_exact(&mut vec![0; content_length]).unwrap();
384
+
385
+ let mut send = |s: &str| {
386
+ stream.write_all(s.as_bytes()).unwrap();
387
+ stream.flush().unwrap();
388
+ };
389
+ send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n");
390
+ for _ in 0..num_tokens {
391
+ send("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n");
392
+ }
393
+ std::thread::sleep(usage_delay);
394
+ send(&format!(
395
+ "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}}}}}\n\n"
396
+ ));
397
+ send("data: [DONE]\n\n");
398
+ });
399
+ format!("http://{addr}/v1/chat/completions")
400
+ }
401
+
402
+ /// The trailing usage chunk carries no token, so it must not extend the
403
+ /// request's end: `latency - ttft` has to equal `sum(itl)`, otherwise
404
+ /// E2EL and TPOT absorb whatever gap precedes the usage chunk. Mirrors the
405
+ /// Python fix in vllm-project/vllm#55508.
406
+ #[tokio::test]
407
+ async fn test_usage_chunk_does_not_extend_latency() {
408
+ // The gap before the usage chunk is the quantity under test, so a
409
+ // real delay is needed here; there is no observable event to wait on.
410
+ let usage_delay = std::time::Duration::from_millis(150);
411
+ let input = RequestFuncInput {
412
+ api_url: spawn_sse_server(3, usage_delay),
413
+ model: "test-model".to_string(),
414
+ output_len: 3,
415
+ ..Default::default()
416
+ };
417
+
418
+ let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap();
419
+
420
+ assert!(output.success, "{}", output.error);
421
+ assert_eq!(output.output_tokens, 3);
422
+ assert_eq!(output.itl.len(), 2);
423
+ let decode_time = output.latency - output.ttft;
424
+ let itl_sum: f64 = output.itl.iter().sum();
425
+ assert!(
426
+ (decode_time - itl_sum).abs() < 1e-6,
427
+ "latency - ttft ({decode_time}s) != sum(itl) ({itl_sum}s)"
428
+ );
429
+ }
430
+
431
/// ignore_eos and extra_body must survive the raw-splice path.
432
#[test]
433
fn test_mm_payload_tail_fields() {