@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 74850f9f66

[Rust Frontend] Optimize SSE streaming hot path (#51321)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Reid committed Aug 29, 2026 at 04:35 UTC 74850f9f669085c07353a5673b1fb2f352480cda
8 files changed +204 -196
rust/Cargo.toml
+1 -1
@@ -90,7 +90,7 @@ rlimit = "0.11.0"
90 rmp-serde = "1.3.1"
91 rmpv = { version = "1.3.1", features = ["with-serde"] }
92 rustc-hash = "1.1.0"
93 -serde = { version = "1.0.228", features = ["derive"] }
93 +serde = { version = "1.0.228", features = ["derive", "rc"] }
94 serde-json-fmt = "0.1.0"
95 serde_default = "0.2.0"
96 serde_json = { version = "1.0.145", features = ["preserve_order"] }
rust/src/server/src/routes/inference/generate.rs
+9 -7
@@ -395,16 +395,18 @@ async fn generate_sse_stream(
395 }
396
397 fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
398 - let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
399 - trace!(payload, "generate emitting chunk");
400 - Event::default().data(payload)
398 + trace!(?chunk, "generate emitting chunk");
399 + Event::default()
400 + .json_data(chunk)
401 + .expect("generate chunk must serialize to JSON")
402 }
403
404 fn to_error_sse_event(error: &ApiError) -> Event {
404 - let payload = serde_json::to_string(&error.to_error_response())
405 - .expect("ErrorResponse must serialize to JSON");
406 - trace!(payload, "generate emitting error");
407 - Event::default().data(payload)
405 + let response = error.to_error_response();
406 + trace!(?response, "generate emitting error");
407 + Event::default()
408 + .json_data(response)
409 + .expect("ErrorResponse must serialize to JSON")
410 }
411
412 fn done_sse_event() -> Event {
rust/src/server/src/routes/openai/chat_completions.rs
+68 -106
@@ -39,7 +39,8 @@ use crate::routes::openai::utils::logprobs::{
39 decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
40 };
41 use crate::routes::openai::utils::types::{
42 - ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
42 + ChatLogProbs, FunctionCallDelta, FunctionCallResponse, StreamResponseEnvelope, ToolCall,
43 + ToolCallDelta, Usage,
44 };
45 use crate::routes::openai::utils::usage::ContinuousUsage;
46 use crate::routes::openai::utils::validated_json::ValidatedJson;
@@ -270,6 +271,12 @@ async fn chat_completion_chunk_stream(
271 }: ResponseOptions,
272 mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
273 ) -> Result<(), ApiError> {
274 + let envelope = Arc::new(StreamResponseEnvelope::new(
275 + request_id,
276 + "chat.completion.chunk",
277 + created,
278 + response_model,
279 + ));
280 let mut saw_tool_calls = false;
281 // `LogprobsDelta` is emitted after all chat events for one decoded update.
282 // If that update contains hidden reasoning, including delimiter-only block
@@ -303,7 +310,7 @@ async fn chat_completion_chunk_stream(
310 prompt_token_ids, ..
311 }) => {
312 continuous_usage.set_prompt_tokens(prompt_token_ids.len());
306 - let mut chunk = start_chunk(&request_id, &response_model, created);
313 + let mut chunk = start_chunk(&envelope);
314 if return_token_ids {
315 chunk.prompt_token_ids = Some(prompt_token_ids.to_vec());
316 }
@@ -311,9 +318,7 @@ async fn chat_completion_chunk_stream(
318 // When echo=true, emit the last assistant message content as a delta chunk.
319 if let Some(echo_text) = &echo {
320 yield_chunk!(block_delta_chunk(
314 - &request_id,
315 - &response_model,
316 - created,
321 + &envelope,
322 AssistantBlockKind::Text,
323 echo_text.clone(),
324 ));
@@ -326,13 +331,7 @@ async fn chat_completion_chunk_stream(
331 if let Some(pending_chunk) = pending_chunk.as_mut() {
332 pending_chunk.push_block_delta(kind, delta);
333 } else {
329 - yield_chunk!(block_delta_chunk(
330 - &request_id,
331 - &response_model,
332 - created,
333 - kind,
334 - delta,
335 - ));
334 + yield_chunk!(block_delta_chunk(&envelope, kind, delta));
335 }
336 } else {
337 suppress_current_update_metadata = true;
@@ -368,18 +367,11 @@ async fn chat_completion_chunk_stream(
367 if let Some(pending_chunk) = pending_chunk.as_mut() {
368 pending_chunk.logprobs = openai_logprobs;
369 pending_chunk.token_ids = openai_token_ids;
371 - if let Some(chunk) =
372 - pending_chunk.take_chunk(&request_id, &response_model, created)
373 - {
370 + if let Some(chunk) = pending_chunk.take_chunk(&envelope) {
371 yield_chunk!(chunk);
372 }
373 } else if let Some(logprobs) = openai_logprobs {
377 - yield_chunk!(logprobs_only_chunk(
378 - &request_id,
379 - &response_model,
380 - created,
381 - logprobs,
382 - ));
374 + yield_chunk!(logprobs_only_chunk(&envelope, logprobs));
375 }
376 }
377 Ok(ChatEvent::BlockStart { kind, .. }) => {
@@ -407,14 +399,7 @@ async fn chat_completion_chunk_stream(
399 if let Some(pending_chunk) = pending_chunk.as_mut() {
400 pending_chunk.push_tool_call_start(tool_index, id, name);
401 } else {
410 - yield_chunk!(tool_call_start_chunk(
411 - &request_id,
412 - &response_model,
413 - created,
414 - tool_index,
415 - id,
416 - name,
417 - ));
402 + yield_chunk!(tool_call_start_chunk(&envelope, tool_index, id, name));
403 }
404 }
405 Ok(ChatEvent::ToolCallArgumentsDelta { index, delta }) => {
@@ -422,13 +407,7 @@ async fn chat_completion_chunk_stream(
407 if let Some(pending_chunk) = pending_chunk.as_mut() {
408 pending_chunk.push_tool_call_arguments(tool_index, delta);
409 } else {
425 - yield_chunk!(tool_call_arguments_chunk(
426 - &request_id,
427 - &response_model,
428 - created,
429 - tool_index,
430 - delta,
431 - ));
410 + yield_chunk!(tool_call_arguments_chunk(&envelope, tool_index, delta));
411 }
412 }
413 Ok(ChatEvent::ToolCallEnd { .. }) => {
@@ -442,7 +421,7 @@ async fn chat_completion_chunk_stream(
421 if enable_log_requests {
422 info!(
423 stream = true,
445 - model = %response_model,
424 + model = %envelope.model(),
425 prompt_tokens = final_usage.prompt_token_count,
426 output_tokens = final_usage.output_token_count,
427 finish_reason = finish_reason.as_str(),
@@ -456,16 +435,13 @@ async fn chat_completion_chunk_stream(
435 );
436
437 if let Some(pending_chunk) = pending_chunk.as_mut()
459 - && let Some(chunk) =
460 - pending_chunk.take_chunk(&request_id, &response_model, created)
438 + && let Some(chunk) = pending_chunk.take_chunk(&envelope)
439 {
440 yield_chunk!(chunk);
441 }
442
443 match final_chunk(
466 - &request_id,
467 - &response_model,
468 - created,
444 + &envelope,
445 finish_reason,
446 saw_tool_calls && !is_named_tool_choice,
447 ) {
@@ -481,9 +457,7 @@ async fn chat_completion_chunk_stream(
457
458 if include_usage {
459 y.yield_ok(usage_chunk(
484 - &request_id,
485 - &response_model,
486 - created,
460 + &envelope,
461 Usage::from_token_usage(final_usage, enable_prompt_tokens_details),
462 ))
463 .await;
@@ -504,12 +478,10 @@ async fn chat_completion_chunk_stream(
478 }
479
480 fn usage_chunk(
507 - request_id: &str,
508 - response_model: &str,
509 - created: u64,
481 + envelope: &Arc<StreamResponseEnvelope>,
482 usage: Usage,
483 ) -> ChatCompletionStreamResponse {
512 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
484 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
485 chunk.usage = Some(usage);
486 chunk
487 }
@@ -588,9 +560,7 @@ impl PendingChatChunk {
560 /// necessarily with a visible/chat-semantic delta.
561 fn take_chunk(
562 &mut self,
591 - request_id: &str,
592 - response_model: &str,
593 - created: u64,
563 + envelope: &Arc<StreamResponseEnvelope>,
564 ) -> Option<ChatCompletionStreamResponse> {
565 let has_delta = self.delta.content.is_some()
566 || self.delta.reasoning.is_some()
@@ -601,7 +571,7 @@ impl PendingChatChunk {
571 return None;
572 }
573
604 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
574 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
575 chunk.choices.push(ChatCompletionStreamChoice {
576 delta: self.take_delta(),
577 logprobs,
@@ -659,18 +629,19 @@ async fn chat_completion_sse_stream(
629
630 /// Serialize one OpenAI chunk payload into one SSE `data:` event.
631 fn to_sse_event(chunk: &ChatCompletionStreamResponse) -> Event {
662 - let payload =
663 - serde_json::to_string(chunk).expect("ChatCompletionStreamResponse must serialize to JSON");
664 - trace!(payload, "chat completion emitting chunk");
665 - Event::default().data(payload)
632 + trace!(?chunk, "chat completion emitting chunk");
633 + Event::default()
634 + .json_data(chunk)
635 + .expect("ChatCompletionStreamResponse must serialize to JSON")
636 }
637
638 /// Serialize one OpenAI error payload into one SSE `data:` event.
639 fn to_error_sse_event(error: &ApiError) -> Event {
670 - let payload = serde_json::to_string(&error.to_error_response())
671 - .expect("ErrorResponse must serialize to JSON");
672 - trace!(payload, "chat completion emitting error");
673 - Event::default().data(payload)
640 + let response = error.to_error_response();
641 + trace!(?response, "chat completion emitting error");
642 + Event::default()
643 + .json_data(response)
644 + .expect("ErrorResponse must serialize to JSON")
645 }
646
647 /// Build the terminal OpenAI SSE sentinel event.
@@ -681,12 +652,8 @@ fn done_sse_event() -> Event {
652
653 /// Build the initial assistant-role SSE chunk required by the OpenAI streaming
654 /// protocol.
684 -fn start_chunk(
685 - request_id: &str,
686 - response_model: &str,
687 - created: u64,
688 -) -> ChatCompletionStreamResponse {
689 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
655 +fn start_chunk(envelope: &Arc<StreamResponseEnvelope>) -> ChatCompletionStreamResponse {
656 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
657 chunk.choices.push(ChatCompletionStreamChoice {
658 delta: ChatMessageDelta {
659 role: Some(AssistantRole),
@@ -699,9 +666,7 @@ fn start_chunk(
666
667 /// Build one content-delta SSE chunk from one internal assistant block delta.
668 fn block_delta_chunk(
702 - request_id: &str,
703 - response_model: &str,
704 - created: u64,
669 + envelope: &Arc<StreamResponseEnvelope>,
670 kind: AssistantBlockKind,
671 delta: String,
672 ) -> ChatCompletionStreamResponse {
@@ -719,7 +684,7 @@ fn block_delta_chunk(
684 }
685 };
686
722 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
687 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
688 chunk.choices.push(ChatCompletionStreamChoice {
689 delta,
690 ..Default::default()
@@ -728,14 +693,12 @@ fn block_delta_chunk(
693 }
694
695 fn tool_call_start_chunk(
731 - request_id: &str,
732 - response_model: &str,
733 - created: u64,
696 + envelope: &Arc<StreamResponseEnvelope>,
697 tool_index: u32,
698 id: String,
699 name: String,
700 ) -> ChatCompletionStreamResponse {
738 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
701 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
702 chunk.choices.push(ChatCompletionStreamChoice {
703 delta: ChatMessageDelta {
704 tool_calls: Some(vec![ToolCallDelta {
@@ -755,13 +718,11 @@ fn tool_call_start_chunk(
718 }
719
720 fn tool_call_arguments_chunk(
758 - request_id: &str,
759 - response_model: &str,
760 - created: u64,
721 + envelope: &Arc<StreamResponseEnvelope>,
722 tool_index: u32,
723 delta: String,
724 ) -> ChatCompletionStreamResponse {
764 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
725 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
726 chunk.choices.push(ChatCompletionStreamChoice {
727 delta: ChatMessageDelta {
728 tool_calls: Some(vec![ToolCallDelta {
@@ -781,12 +742,10 @@ fn tool_call_arguments_chunk(
742 }
743
744 fn logprobs_only_chunk(
784 - request_id: &str,
785 - response_model: &str,
786 - created: u64,
745 + envelope: &Arc<StreamResponseEnvelope>,
746 logprobs: ChatLogProbs,
747 ) -> ChatCompletionStreamResponse {
789 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
748 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
749 chunk.choices.push(ChatCompletionStreamChoice {
750 logprobs: Some(logprobs),
751 ..Default::default()
@@ -796,9 +755,7 @@ fn logprobs_only_chunk(
755
756 /// Build the terminal SSE chunk carrying the OpenAI finish reason.
757 fn final_chunk(
799 - request_id: &str,
800 - response_model: &str,
801 - created: u64,
758 + envelope: &Arc<StreamResponseEnvelope>,
759 finish_reason: FinishReason,
760 use_tool_calls_finish_reason: bool,
761 ) -> Result<ChatCompletionStreamResponse, ApiError> {
@@ -811,7 +768,7 @@ fn final_chunk(
768 "chat stream finished"
769 );
770
814 - let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created);
771 + let mut chunk = ChatCompletionStreamResponse::new(envelope);
772 chunk.choices.push(ChatCompletionStreamChoice {
773 finish_reason: Some(finish_reason.to_string()),
774 stop_reason,
@@ -844,6 +801,8 @@ fn stop_reason_to_json(stop_reason: &StopReason) -> Value {
801
802 #[cfg(test)]
803 mod tests {
804 + use std::sync::Arc;
805 +
806 use futures::{StreamExt as _, stream};
807 use serde_json::json;
808 use vllm_chat::{
@@ -853,19 +812,23 @@ mod tests {
812 use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
813
814 use super::{
856 - ApiServerOptions, ResponseOptions, block_delta_chunk, chat_completion_chunk_stream,
857 - final_chunk,
815 + ApiServerOptions, ResponseOptions, StreamResponseEnvelope, block_delta_chunk,
816 + chat_completion_chunk_stream, final_chunk,
817 };
818
819 + fn stream_envelope() -> Arc<StreamResponseEnvelope> {
820 + Arc::new(StreamResponseEnvelope::new(
821 + "chatcmpl-1".to_string(),
822 + "chat.completion.chunk",
823 + 1,
824 + "model".to_string(),
825 + ))
826 + }
827 +
828 #[test]
829 fn text_chunk_uses_content_only_delta() {
862 - let chunk = block_delta_chunk(
863 - "chatcmpl-1",
864 - "model",
865 - 1,
866 - AssistantBlockKind::Text,
867 - "hello".to_string(),
868 - );
830 + let envelope = stream_envelope();
831 + let chunk = block_delta_chunk(&envelope, AssistantBlockKind::Text, "hello".to_string());
832 assert_eq!(chunk.choices[0].delta.role, None);
833 assert_eq!(chunk.choices[0].delta.content.as_deref(), Some("hello"));
834 assert_eq!(chunk.choices[0].delta.reasoning, None);
@@ -873,10 +836,9 @@ mod tests {
836
837 #[test]
838 fn reasoning_chunk_uses_reasoning_only_delta() {
839 + let envelope = stream_envelope();
840 let chunk = block_delta_chunk(
877 - "chatcmpl-1",
878 - "model",
879 - 1,
841 + &envelope,
842 AssistantBlockKind::Reasoning,
843 "thinking".to_string(),
844 );
@@ -890,10 +852,9 @@ mod tests {
852
853 #[test]
854 fn final_chunk_maps_stop_finish_reason_and_stop_reason() {
855 + let envelope = stream_envelope();
856 let chunk = final_chunk(
894 - "chatcmpl-1",
895 - "model",
896 - 1,
857 + &envelope,
858 FinishReason::Stop(Some(StopReason::Text("stop".to_string()))),
859 false,
860 )
@@ -905,7 +866,7 @@ mod tests {
866
867 #[test]
868 fn final_chunk_maps_length_finish_reason() {
908 - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Length, false)
869 + let chunk = final_chunk(&stream_envelope(), FinishReason::Length, false)
870 .expect("finish reason is valid");
871
872 assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
@@ -914,7 +875,7 @@ mod tests {
875
876 #[test]
877 fn final_chunk_maps_abort_finish_reason() {
917 - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Abort, false)
878 + let chunk = final_chunk(&stream_envelope(), FinishReason::Abort, false)
879 .expect("abort is a valid finish reason");
880
881 assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
@@ -923,12 +884,12 @@ mod tests {
884
885 #[test]
886 fn final_chunk_rejects_error_finish_reason() {
926 - assert!(final_chunk("chatcmpl-1", "model", 1, FinishReason::Error, false).is_err());
887 + assert!(final_chunk(&stream_envelope(), FinishReason::Error, false).is_err());
888 }
889
890 #[test]
891 fn final_chunk_maps_stop_to_tool_calls_when_tool_calls_were_streamed() {
931 - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::stop_eos(), true)
892 + let chunk = final_chunk(&stream_envelope(), FinishReason::stop_eos(), true)
893 .expect("finish reason is valid");
894
895 assert_eq!(
@@ -1002,6 +963,7 @@ mod tests {
963 .expect("stream chunks");
964
965 assert_eq!(chunks.len(), 4);
966 + assert!(chunks.windows(2).all(|pair| Arc::ptr_eq(&pair[0].envelope, &pair[1].envelope)));
967 assert_eq!(chunks[1].choices[0].delta.content.as_deref(), Some("hi"));
968 let logprobs = chunks[1].choices[0].logprobs.as_ref().expect("logprobs");
969 let content = logprobs.content.as_ref().expect("logprobs content");
rust/src/server/src/routes/openai/chat_completions/types.rs
+8 -12
@@ -3,6 +3,7 @@
3
4 use std::collections::HashMap;
5 use std::fmt;
6 +use std::sync::Arc;
7
8 use serde::{Deserialize, Serialize};
9 use serde_json::Value;
@@ -13,9 +14,9 @@ use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;
14
15 use crate::routes::openai::utils::structured_outputs::ResponseFormat;
16 use crate::routes::openai::utils::types::{
16 - ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, StreamOptions, StringOrArray, Tool,
17 - ToolCall, ToolCallDelta, ToolChoice, Usage, default_true, deserialize_request_top_k,
18 - validate_messages, validate_stop, validate_top_p_value,
17 + ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope,
18 + StringOrArray, Tool, ToolCall, ToolCallDelta, ToolChoice, Usage, default_true,
19 + deserialize_request_top_k, validate_messages, validate_stop, validate_top_p_value,
20 };
21
22 /// vLLM-compatible request type for the Chat Completions API.
@@ -383,10 +384,8 @@ pub(super) struct ChatCompletionMessage {
384 #[serde_with::skip_serializing_none]
385 #[derive(Debug, Clone, Serialize)]
386 pub(super) struct ChatCompletionStreamResponse {
386 - pub id: String,
387 - pub object: String,
388 - pub created: u64,
389 - pub model: String,
387 + #[serde(flatten)]
388 + pub envelope: Arc<StreamResponseEnvelope>,
389 pub choices: Vec<ChatCompletionStreamChoice>,
390 pub usage: Option<Usage>,
391 pub prompt_token_ids: Option<Vec<u32>>,
@@ -394,12 +393,9 @@ pub(super) struct ChatCompletionStreamResponse {
393
394 impl ChatCompletionStreamResponse {
395 /// Create a stream response with the standard envelope fields pre-filled.
397 - pub fn new(id: &str, model: &str, created: u64) -> Self {
396 + pub fn new(envelope: &Arc<StreamResponseEnvelope>) -> Self {
397 Self {
399 - id: id.to_string(),
400 - object: "chat.completion.chunk".to_string(),
401 - created,
402 - model: model.to_string(),
398 + envelope: Arc::clone(envelope),
399 choices: Vec::new(),
400 usage: None,
401 prompt_token_ids: None,
rust/src/server/src/routes/openai/completions.rs
+62 -59
@@ -33,7 +33,7 @@ use super::utils::logprobs::{
33 collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
34 prompt_logprobs_to_maps, text_len,
35 };
36 -use super::utils::types::Usage;
36 +use super::utils::types::{StreamResponseEnvelope, Usage};
37 use crate::config::ApiServerOptions;
38 use crate::error::{ApiError, bail_server_error, server_error, text_submit_error};
39 use crate::lora::LoraModelResolution;
@@ -257,6 +257,12 @@ async fn completion_chunk_stream(
257 mut y: TryYielder<CompletionSseChunk, ApiError>,
258 ) -> Result<(), ApiError> {
259 pin_mut!(stream);
260 + let envelope = Arc::new(StreamResponseEnvelope::new(
261 + request_id,
262 + "text_completion",
263 + created,
264 + response_model,
265 + ));
266 let mut visible_text_len = 0_u32;
267 let mut first_chunk = true;
268 let mut continuous_usage = ContinuousUsage::default();
@@ -292,13 +298,7 @@ async fn completion_chunk_stream(
298 } else {
299 None
300 };
295 - let mut chunk = delta_chunk(
296 - &request_id,
297 - &response_model,
298 - created,
299 - prompt.clone(),
300 - logprobs,
301 - );
301 + let mut chunk = delta_chunk(&envelope, prompt.clone(), logprobs);
302 if return_token_ids && first_chunk {
303 if let Some(choice) = chunk.choices.first_mut() {
304 choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
@@ -308,8 +308,7 @@ async fn completion_chunk_stream(
308 yield_chunk!(chunk);
309 } else if return_token_ids {
310 // Emit a chunk with prompt_token_ids in the first streaming response
311 - let mut chunk =
312 - delta_chunk(&request_id, &response_model, created, String::new(), None);
311 + let mut chunk = delta_chunk(&envelope, String::new(), None);
312 if let Some(choice) = chunk.choices.first_mut() {
313 choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
314 }
@@ -331,7 +330,7 @@ async fn completion_chunk_stream(
330 if enable_log_requests {
331 info!(
332 stream = true,
334 - model = %response_model,
333 + model = %envelope.model(),
334 prompt_tokens = finished.usage.prompt_token_count,
335 output_tokens = finished.usage.output_token_count,
336 finish_reason = finished.finish_reason.as_str(),
@@ -342,19 +341,12 @@ async fn completion_chunk_stream(
341 finished.usage.prompt_token_count,
342 finished.usage.output_token_count,
343 );
345 - let final_chunk = final_chunk(
346 - &request_id,
347 - &response_model,
348 - created,
349 - finished.finish_reason,
350 - )?;
344 + let final_chunk = final_chunk(&envelope, finished.finish_reason)?;
345 yield_chunk!(final_chunk);
346
347 if include_usage {
348 y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
355 - &request_id,
356 - &response_model,
357 - created,
349 + &envelope,
350 Usage::from_token_usage(
351 finished.usage,
352 enable_prompt_tokens_details,
@@ -380,7 +372,7 @@ async fn completion_chunk_stream(
372 } else {
373 None
374 };
383 - let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs);
375 + let mut chunk = delta_chunk(&envelope, delta, logprobs);
376 let delta_token_count = token_ids.len();
377 continuous_usage.add_output_tokens(delta_token_count);
378 if return_token_ids && let Some(choice) = chunk.choices.first_mut() {
@@ -393,7 +385,7 @@ async fn completion_chunk_stream(
385 if enable_log_requests {
386 info!(
387 stream = true,
396 - model = %response_model,
388 + model = %envelope.model(),
389 prompt_tokens = finished.usage.prompt_token_count,
390 output_tokens = finished.usage.output_token_count,
391 finish_reason = finished.finish_reason.as_str(),
@@ -404,19 +396,12 @@ async fn completion_chunk_stream(
396 finished.usage.prompt_token_count,
397 finished.usage.output_token_count,
398 );
407 - let final_chunk = final_chunk(
408 - &request_id,
409 - &response_model,
410 - created,
411 - finished.finish_reason,
412 - )?;
399 + let final_chunk = final_chunk(&envelope, finished.finish_reason)?;
400 yield_chunk!(final_chunk);
401
402 if include_usage {
403 y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
417 - &request_id,
418 - &response_model,
419 - created,
404 + &envelope,
405 Usage::from_token_usage(finished.usage, enable_prompt_tokens_details),
406 )))
407 .await;
@@ -436,13 +421,11 @@ async fn completion_chunk_stream(
421 }
422
423 fn delta_chunk(
439 - request_id: &str,
440 - response_model: &str,
441 - created: u64,
424 + envelope: &Arc<StreamResponseEnvelope>,
425 text: String,
426 logprobs: Option<LogProbs>,
427 ) -> CompletionStreamResponse {
445 - let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
428 + let mut chunk = CompletionStreamResponse::new(envelope);
429 chunk.choices.push(CompletionStreamChoice {
430 text,
431 logprobs,
@@ -452,15 +435,13 @@ fn delta_chunk(
435 }
436
437 fn final_chunk(
455 - request_id: &str,
456 - response_model: &str,
457 - created: u64,
438 + envelope: &Arc<StreamResponseEnvelope>,
439 finish_reason: FinishReason,
440 ) -> Result<CompletionStreamResponse, ApiError> {
441 let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
442 let finish_reason = completion_finish_reason_to_openai(&finish_reason)?;
443
463 - let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
444 + let mut chunk = CompletionStreamResponse::new(envelope);
445 chunk.choices.push(CompletionStreamChoice {
446 finish_reason: Some(finish_reason.to_string()),
447 stop_reason,
@@ -517,13 +498,8 @@ fn prompt_only_logprobs_to_openai(
498 ))
499 }
500
520 -fn usage_chunk(
521 - request_id: &str,
522 - response_model: &str,
523 - created: u64,
524 - usage: Usage,
525 -) -> CompletionStreamResponse {
526 - let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);
501 +fn usage_chunk(envelope: &Arc<StreamResponseEnvelope>, usage: Usage) -> CompletionStreamResponse {
502 + let mut chunk = CompletionStreamResponse::new(envelope);
503 chunk.usage = Some(usage);
504 chunk
505 }
@@ -556,17 +532,19 @@ async fn completion_sse_stream(
532
533 /// Serialize one OpenAI chunk payload into one SSE `data:` event.
534 fn to_sse_event(chunk: &CompletionSseChunk) -> Event {
559 - let payload = serde_json::to_string(chunk).expect("completion chunk must serialize to JSON");
560 - trace!(payload, "completion emitting chunk");
561 - Event::default().data(payload)
535 + trace!(?chunk, "completion emitting chunk");
536 + Event::default()
537 + .json_data(chunk)
538 + .expect("completion chunk must serialize to JSON")
539 }
540
541 /// Serialize one OpenAI error payload into one SSE `data:` event.
542 fn to_error_sse_event(error: &ApiError) -> Event {
566 - let payload = serde_json::to_string(&error.to_error_response())
567 - .expect("ErrorResponse must serialize to JSON");
568 - trace!(payload, "completion emitting error");
569 - Event::default().data(payload)
543 + let response = error.to_error_response();
544 + trace!(?response, "completion emitting error");
545 + Event::default()
546 + .json_data(response)
547 + .expect("ErrorResponse must serialize to JSON")
548 }
549
550 /// Build the terminal OpenAI SSE sentinel event.
@@ -577,6 +555,8 @@ fn done_sse_event() -> Event {
555
556 #[cfg(test)]
557 mod tests {
558 + use std::sync::Arc;
559 +
560 use futures::{StreamExt as _, stream};
561 use itertools::Itertools as _;
562 use vllm_engine_core_client::protocol::output::StopReason;
@@ -586,13 +566,29 @@ mod tests {
566 };
567
568 use super::{
589 - ApiServerOptions, CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk,
569 + ApiServerOptions, CompletionSseChunk, CompletionStreamResponse, ResponseOptions,
570 + StreamResponseEnvelope, completion_chunk_stream, final_chunk,
571 };
572
573 + fn stream_envelope() -> Arc<StreamResponseEnvelope> {
574 + Arc::new(StreamResponseEnvelope::new(
575 + "cmpl-1".to_string(),
576 + "text_completion",
577 + 1,
578 + "model".to_string(),
579 + ))
580 + }
581 +
582 + fn chunk_response(chunk: &CompletionSseChunk) -> &CompletionStreamResponse {
583 + match chunk {
584 + CompletionSseChunk::Chunk(response) | CompletionSseChunk::Usage(response) => response,
585 + }
586 + }
587 +
588 #[test]
589 fn final_chunk_maps_stop_finish_reason() {
594 - let chunk = final_chunk("cmpl-1", "model", 1, FinishReason::stop_eos())
595 - .expect("finish reason valid");
590 + let chunk =
591 + final_chunk(&stream_envelope(), FinishReason::stop_eos()).expect("finish reason valid");
592 assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("stop"));
593 assert_eq!(chunk.choices[0].text, "");
594 }
@@ -600,20 +596,20 @@ mod tests {
596 #[test]
597 fn final_chunk_maps_length_finish_reason() {
598 let chunk =
603 - final_chunk("cmpl-1", "model", 1, FinishReason::Length).expect("finish reason valid");
599 + final_chunk(&stream_envelope(), FinishReason::Length).expect("finish reason valid");
600 assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
601 }
602
603 #[test]
604 fn final_chunk_maps_abort_finish_reason() {
605 let chunk =
610 - final_chunk("cmpl-1", "model", 1, FinishReason::Abort).expect("finish reason valid");
606 + final_chunk(&stream_envelope(), FinishReason::Abort).expect("finish reason valid");
607 assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
608 }
609
610 #[test]
611 fn final_chunk_rejects_error_finish_reason() {
616 - assert!(final_chunk("cmpl-1", "model", 1, FinishReason::Error).is_err());
612 + assert!(final_chunk(&stream_envelope(), FinishReason::Error).is_err());
613 }
614
615 #[tokio::test]
@@ -702,6 +698,13 @@ mod tests {
698
699 let chunks: Vec<_> = chunks.into_iter().try_collect().expect("stream should succeed");
700
701 + assert!(chunks.windows(2).all(|pair| {
702 + Arc::ptr_eq(
703 + &chunk_response(&pair[0]).envelope,
704 + &chunk_response(&pair[1]).envelope,
705 + )
706 + }));
707 +
708 match &chunks[0] {
709 CompletionSseChunk::Chunk(chunk) => {
710 assert_eq!(chunk.choices[0].text, "h");
rust/src/server/src/routes/openai/completions/types.rs
+7 -11
@@ -2,6 +2,7 @@
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::collections::HashMap;
5 +use std::sync::Arc;
6
7 use serde::{Deserialize, Serialize};
8 use serde_json::{Map, Value};
@@ -10,8 +11,8 @@ use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;
11 use vllm_text::Prompt;
12
13 use crate::routes::openai::utils::types::{
13 - LogProbs, Normalizable, PromptLogprobs, StreamOptions, StringOrArray, Usage, default_true,
14 - deserialize_request_top_k, validate_stop,
14 + LogProbs, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope, StringOrArray,
15 + Usage, default_true, deserialize_request_top_k, validate_stop,
16 };
17
18 /// Serde default for `CompletionRequest::max_tokens`, matching the Python vLLM
@@ -241,22 +242,17 @@ pub(super) struct CompletionChoice {
242 #[serde_with::skip_serializing_none]
243 #[derive(Debug, Clone, Serialize)]
244 pub(super) struct CompletionStreamResponse {
244 - pub id: String,
245 - pub object: String,
246 - pub created: u64,
247 - pub model: String,
245 + #[serde(flatten)]
246 + pub envelope: Arc<StreamResponseEnvelope>,
247 pub choices: Vec<CompletionStreamChoice>,
248 pub usage: Option<Usage>,
249 }
250
251 impl CompletionStreamResponse {
252 /// Create a stream response with the standard envelope fields pre-filled.
254 - pub fn new(id: &str, model: &str, created: u64) -> Self {
253 + pub fn new(envelope: &Arc<StreamResponseEnvelope>) -> Self {
254 Self {
256 - id: id.to_string(),
257 - object: "text_completion".to_string(),
258 - created,
259 - model: model.to_string(),
255 + envelope: Arc::clone(envelope),
256 choices: Vec::new(),
257 usage: None,
258 }
rust/src/server/src/routes/openai/utils/types.rs
+23
@@ -201,6 +201,29 @@ pub struct InputAudio {
201 // Streaming
202 // ============================================================================
203
204 +#[derive(Debug, Serialize)]
205 +pub(crate) struct StreamResponseEnvelope {
206 + id: String,
207 + object: &'static str,
208 + created: u64,
209 + model: String,
210 +}
211 +
212 +impl StreamResponseEnvelope {
213 + pub(crate) fn new(id: String, object: &'static str, created: u64, model: String) -> Self {
214 + Self {
215 + id,
216 + object,
217 + created,
218 + model,
219 + }
220 + }
221 +
222 + pub(crate) fn model(&self) -> &str {
223 + &self.model
224 + }
225 +}
226 +
227 /// Mirrors the Python vLLM `StreamOptions` class.
228 #[derive(Debug, Clone, Deserialize, Serialize)]
229 pub struct StreamOptions {
rust/src/server/src/routes/tests.rs
+26
@@ -170,6 +170,22 @@ fn sse_json_payloads(text: &str) -> Vec<serde_json::Value> {
170 .collect()
171 }
172
173 +fn assert_shared_openai_stream_envelope(payloads: &[serde_json::Value], object: &str, model: &str) {
174 + let first = payloads.first().expect("at least one SSE JSON payload");
175 + assert!(first["id"].as_str().is_some_and(|id| !id.is_empty()));
176 + assert!(first["created"].is_u64());
177 + assert_eq!(first["object"], object);
178 + assert_eq!(first["model"], model);
179 +
180 + for payload in payloads {
181 + assert_eq!(payload["id"], first["id"]);
182 + assert_eq!(payload["object"], first["object"]);
183 + assert_eq!(payload["created"], first["created"]);
184 + assert_eq!(payload["model"], first["model"]);
185 + assert!(payload.get("envelope").is_none());
186 + }
187 +}
188 +
189 type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
190
191 fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
@@ -2717,6 +2733,11 @@ async fn happy_path_returns_sse_stream() {
2733 let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2734 let after = METRICS.render().unwrap();
2735
2736 + assert_shared_openai_stream_envelope(
2737 + &sse_json_payloads(&text),
2738 + "chat.completion.chunk",
2739 + "Qwen/Qwen1.5-0.5B-Chat",
2740 + );
2741 assert!(text.contains("\"role\":\"assistant\""), "{text}");
2742 assert!(text.starts_with("data: "), "{text}");
2743 assert_eq!(
@@ -4465,6 +4486,11 @@ async fn completions_happy_path_returns_sse_stream() {
4486 engine_task.await.expect("mock engine task");
4487 let text = String::from_utf8(body.to_vec()).expect("utf8 body");
4488 let payloads = sse_data_payloads(&text);
4489 + assert_shared_openai_stream_envelope(
4490 + &sse_json_payloads(&text),
4491 + "text_completion",
4492 + "Qwen/Qwen1.5-0.5B-Chat",
4493 + );
4494 let usage_index = payloads
4495 .iter()
4496 .position(|payload| payload.contains("\"usage\":"))