[Rust Frontend] Box rarely-set fields of per-token output types (#60588)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

rust/src/chat/src/event.rs

@@ -237,9 +237,11 @@ pub enum ChatEvent {
         usage: ChatTokenUsage,
         finish_reason: FinishReason,
         /// Connector-specific KV transfer parameters for disaggregated serving.
-        kv_transfer_params: Option<serde_json::Value>,
+        /// Boxed, like `ec_transfer_params`, to keep the per-token variants of
+        /// this enum small.
+        kv_transfer_params: Option<Box<serde_json::Value>>,
         /// Connector-specific encoder cache transfer parameters for
         /// disaggregated serving.
-        ec_transfer_params: Option<serde_json::Value>,
+        ec_transfer_params: Option<Box<serde_json::Value>>,
     },
 }

rust/src/chat/src/output/mod.rs

@@ -54,10 +54,12 @@ pub(crate) enum AssistantEvent {
         usage: ChatTokenUsage,
         finish_reason: FinishReason,
         /// Connector-specific KV transfer parameters for disaggregated serving.
-        kv_transfer_params: Option<serde_json::Value>,
+        /// Boxed, like `ec_transfer_params`, to keep the per-token variants of
+        /// this enum small.
+        kv_transfer_params: Option<Box<serde_json::Value>>,
         /// Connector-specific encoder cache transfer parameters for
         /// disaggregated serving.
-        ec_transfer_params: Option<serde_json::Value>,
+        ec_transfer_params: Option<Box<serde_json::Value>>,
     },
 }
 

rust/src/chat/src/output/structured.rs

@@ -150,8 +150,8 @@ impl StructuredEventState {
         &mut self,
         usage: ChatTokenUsage,
         finish_reason: FinishReason,
-        kv_transfer_params: Option<serde_json::Value>,
-        ec_transfer_params: Option<serde_json::Value>,
+        kv_transfer_params: Option<Box<serde_json::Value>>,
+        ec_transfer_params: Option<Box<serde_json::Value>>,
     ) -> Result<Vec<ChatEvent>> {
         let mut events = Vec::new();
         self.close_open_text_block(&mut events);

rust/src/chat/src/stream.rs

@@ -95,8 +95,8 @@ impl ChatEventStream {
                         token_ids,
                         usage,
                         finish_reason,
-                        kv_transfer_params,
-                        ec_transfer_params,
+                        kv_transfer_params: kv_transfer_params.map(|value| *value),
+                        ec_transfer_params: ec_transfer_params.map(|value| *value),
                     });
                 }
                 ChatEvent::ToolCallEnd { call, .. } => {

rust/src/engine-core-client/src/protocol/output.rs

@@ -82,6 +82,10 @@ pub struct EngineCoreEvent {
 ///
 /// Original Python definition:
 /// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
+///
+/// Fields that are set on few outputs (first or final output, or opt-in
+/// features) are boxed: outputs are moved by value through the dispatcher and
+/// the request stream once per engine step, so their inline size matters.
 #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
 pub struct EngineCoreOutput {
     pub request_id: String,
@@ -95,7 +99,7 @@ pub struct EngineCoreOutput {
     #[serde(default)]
     pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
     #[serde(default)]
-    pub pooling_output: Option<OpaqueValue>,
+    pub pooling_output: Option<Box<OpaqueValue>>,
     #[serde(default)]
     pub finish_reason: Option<EngineCoreFinishReason>,
     #[serde(default)]
@@ -103,17 +107,17 @@ pub struct EngineCoreOutput {
     #[serde(default)]
     pub events: Option<Vec<EngineCoreEvent>>,
     #[serde(default)]
-    pub kv_transfer_params: Option<serde_json::Value>,
+    pub kv_transfer_params: Option<Box<serde_json::Value>>,
     #[serde(default)]
-    pub ec_transfer_params: Option<serde_json::Value>,
+    pub ec_transfer_params: Option<Box<serde_json::Value>>,
     #[serde(default)]
-    pub trace_headers: Option<OpaqueValue>,
+    pub trace_headers: Option<Box<OpaqueValue>>,
     /// Breakdown of the scheduled prefill computation, set on the first output
     /// of a newly scheduled prefill and elided for subsequent decode outputs.
     #[serde(default)]
-    pub prefill_stats: Option<PrefillStats>,
+    pub prefill_stats: Option<Box<PrefillStats>>,
     #[serde(default)]
-    pub routed_experts: Option<OpaqueValue>,
+    pub routed_experts: Option<Box<OpaqueValue>>,
     /// Number of NaNs seen in logits. Values above zero indicate corruption.
     #[serde(default)]
     pub num_nans_in_logits: u32,
@@ -133,11 +137,11 @@ pub struct EngineCoreOutput {
     /// Per-request speculative-decoding acceptance metrics, set on the final
     /// output when `--per-request-spec-decode-metrics` is enabled.
     #[serde(default)]
-    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+    pub spec_decode_metrics: Option<Box<RequestSpecDecodeMetrics>>,
     /// Log probabilities of `SamplingParams.prompt_logprob_token_ids`, set on
     /// the first output of a request that asked for them.
     #[serde(default)]
-    pub prompt_token_id_logprobs: Option<WireNdArray>,
+    pub prompt_token_id_logprobs: Option<Box<WireNdArray>>,
 }
 
 /// Raw per-sequence speculative-decoding accumulator.

rust/src/llm/src/output.rs

@@ -152,7 +152,7 @@ pub struct GenerateOutput {
     pub request_id: String,
     /// One-time prompt metadata emitted only on the first output for this
     /// request.
-    pub prompt_info: Option<GeneratePromptInfo>,
+    pub prompt_info: Option<Box<GeneratePromptInfo>>,
     /// Newly produced token IDs for this step.
     pub token_ids: Vec<u32>,
     /// Sample logprobs for the generated positions in this step.
@@ -162,14 +162,14 @@ pub struct GenerateOutput {
     /// Number of prompt tokens served from cache, when reported by prefill stats.
     pub cached_token_count: usize,
     /// Connector-specific KV transfer parameters for disaggregated serving.
-    pub kv_transfer_params: Option<serde_json::Value>,
+    pub kv_transfer_params: Option<Box<serde_json::Value>>,
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
-    pub ec_transfer_params: Option<serde_json::Value>,
+    pub ec_transfer_params: Option<Box<serde_json::Value>>,
     /// Sampling support sets aligned one-to-one with `token_ids`.
     pub sampling_mask: Option<SamplingMask>,
     /// Per-request speculative-decoding metrics, present on terminal outputs.
-    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+    pub spec_decode_metrics: Option<Box<RequestSpecDecodeMetrics>>,
 }
 
 impl GenerateOutput {
@@ -207,10 +207,12 @@ impl GenerateOutput {
     ) -> Self {
         Self {
             request_id: String::new(),
-            prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
-                prompt_token_ids: ids,
-                prompt_logprobs: None,
-                prompt_token_id_logprobs: None,
+            prompt_info: prompt_token_ids.map(|ids| {
+                Box::new(GeneratePromptInfo {
+                    prompt_token_ids: ids,
+                    prompt_logprobs: None,
+                    prompt_token_id_logprobs: None,
+                })
             }),
             token_ids,
             logprobs: None,
@@ -231,7 +233,7 @@ impl GenerateOutput {
 ///   with `finish_reason = Abort` before the stream ends.
 /// - For errors or unexpected engine-side closes, the stream terminates with an error.
 pub struct GenerateOutputStream {
-    pending_prompt_info: Option<GeneratePromptInfo>,
+    pending_prompt_info: Option<Box<GeneratePromptInfo>>,
     raw_stream: EngineCoreOutputStream,
     request_metrics: RequestMetricsTracker,
     /// Removes this request's external→internal tracking edge on drop. Held for
@@ -249,11 +251,11 @@ impl GenerateOutputStream {
         request_guard: RequestGuard,
     ) -> Self {
         Self {
-            pending_prompt_info: Some(GeneratePromptInfo {
+            pending_prompt_info: Some(Box::new(GeneratePromptInfo {
                 prompt_token_ids,
                 prompt_logprobs: None,
                 prompt_token_id_logprobs: None,
-            }),
+            })),
             raw_stream,
             request_metrics,
             _request_guard: request_guard,
@@ -291,7 +293,7 @@ impl Stream for GenerateOutputStream {
         if let Some(info) = &mut self.pending_prompt_info
             && info.prompt_token_id_logprobs.is_none()
         {
-            info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs;
+            info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs.map(|logprobs| *logprobs);
         }
 
         let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
@@ -432,9 +434,9 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                         output_token_count: collected.token_ids.len(),
                         cached_token_count,
                     };
-                    collected.kv_transfer_params = output.kv_transfer_params;
-                    collected.ec_transfer_params = output.ec_transfer_params;
-                    collected.spec_decode_metrics = output.spec_decode_metrics;
+                    collected.kv_transfer_params = output.kv_transfer_params.map(|value| *value);
+                    collected.ec_transfer_params = output.ec_transfer_params.map(|value| *value);
+                    collected.spec_decode_metrics = output.spec_decode_metrics.map(|value| *value);
                     if let Some(mask) = collected.sampling_mask.as_ref()
                         && mask.rows.len() != collected.token_ids.len()
                     {

rust/src/llm/src/request_metrics.rs

@@ -404,14 +404,14 @@ mod tests {
                         timestamp: 9.0,
                     },
                 ]),
-                prefill_stats: Some(PrefillStats {
+                prefill_stats: Some(Box::new(PrefillStats {
                     num_prompt_tokens: 64,
                     num_computed_tokens: 60,
                     num_cached_tokens: 4,
                     num_local_cached_tokens: 4,
                     num_external_cached_tokens: 0,
                     ..Default::default()
-                }),
+                })),
                 ..Default::default()
             },
         );

rust/src/llm/tests/generate.rs

@@ -92,8 +92,8 @@ fn request_output_with_logprobs_and_kv(
         new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct),
         new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct),
         finish_reason,
-        kv_transfer_params,
-        ec_transfer_params,
+        kv_transfer_params: kv_transfer_params.map(Box::new),
+        ec_transfer_params: ec_transfer_params.map(Box::new),
         ..Default::default()
     }
 }
@@ -280,11 +280,11 @@ async fn generate_streams_outputs() {
     assert_eq!(first.request_id, internal_id);
     assert_eq!(
         first.prompt_info,
-        Some(GeneratePromptInfo {
+        Some(Box::new(GeneratePromptInfo {
             prompt_token_ids: vec![11, 22].into(),
             prompt_logprobs: Some(prompt_logprobs()),
             prompt_token_id_logprobs: None,
-        })
+        }))
     );
     assert_eq!(first.token_ids, vec![1, 2]);
     assert_eq!(
@@ -330,12 +330,12 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
                     RequestBatchOutputs {
                         outputs: vec![
                             EngineCoreOutput {
-                                prefill_stats: Some(PrefillStats {
+                                prefill_stats: Some(Box::new(PrefillStats {
                                     num_prompt_tokens: 2,
                                     num_cached_tokens: 1,
                                     num_local_cached_tokens: 1,
                                     ..Default::default()
-                                }),
+                                })),
                                 new_sampling_mask: Some(MaybeWireSamplingMask::Direct(
                                     SamplingMask {
                                         rows: vec![vec![1, 33, 99]],
@@ -409,11 +409,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
 async fn collect_output_rejects_partial_sampling_mask() {
     let output = vllm_llm::GenerateOutput {
         request_id: "req-partial-mask".to_string(),
-        prompt_info: Some(GeneratePromptInfo {
+        prompt_info: Some(Box::new(GeneratePromptInfo {
             prompt_token_ids: Arc::from([11_u32, 22]),
             prompt_logprobs: None,
             prompt_token_id_logprobs: None,
-        }),
+        })),
         token_ids: vec![33, 44],
         logprobs: None,
         finish_reason: Some(FinishReason::Length),
@@ -763,11 +763,11 @@ async fn generate_records_request_metrics_in_prometheus_output() {
                         engine_index: 4,
                         timestamp: 10.0,
                         outputs: vec![EngineCoreOutput {
-                            prefill_stats: Some(PrefillStats {
+                            prefill_stats: Some(Box::new(PrefillStats {
                                 num_prompt_tokens: 2,
                                 num_computed_tokens: 2,
                                 ..Default::default()
-                            }),
+                            })),
                             ..request_output_with_events(
                                 &request.request_id,
                                 vec![1],

rust/src/mock-engine/src/engine.rs

@@ -201,7 +201,7 @@ impl ActiveRequest {
 
         let finished = self.generated >= self.max_tokens;
         EngineCoreOutput {
-            prefill_stats,
+            prefill_stats: prefill_stats.map(Box::new),
             ..request_output(
                 self.request_id.clone(),
                 new_token_ids,

rust/src/mock-engine/src/tests.rs

@@ -123,11 +123,11 @@ async fn chunk_size_one_outputs_one_token_per_update() {
     assert_eq!(first.finish_reason, None);
     assert_eq!(
         first.prefill_stats,
-        Some(PrefillStats {
+        Some(Box::new(PrefillStats {
             num_prompt_tokens: 3,
             num_computed_tokens: 3,
             ..Default::default()
-        })
+        }))
     );
     let second = stream.next().await.expect("second").expect("second ok");
     assert_eq!(second.new_token_ids.len(), 1);

rust/src/server/src/grpc/convert.rs

@@ -406,8 +406,8 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result<pb::FinishIn
         num_output_tokens: finished.usage.output_token_count as u32,
         finish_reason,
         stop_reason,
-        kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
-        ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
+        kv_transfer_params: finished.kv_transfer_params.as_deref().and_then(json_to_proto_struct),
+        ec_transfer_params: finished.ec_transfer_params.as_deref().and_then(json_to_proto_struct),
         num_cached_tokens: Some(finished.usage.cached_token_count as u32),
     })
 }

rust/src/server/src/grpc/inference.rs

@@ -255,8 +255,8 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
         let finish_info = vllm_text::Finished {
             usage: collected.usage,
             finish_reason: collected.finish_reason,
-            kv_transfer_params: collected.kv_transfer_params,
-            ec_transfer_params: collected.ec_transfer_params,
+            kv_transfer_params: collected.kv_transfer_params.map(Box::new),
+            ec_transfer_params: collected.ec_transfer_params.map(Box::new),
             sampling_mask: collected.sampling_mask,
         };
 

rust/src/server/src/routes/inference/generate.rs

@@ -176,7 +176,7 @@ async fn generate_chunk_stream(
         match next {
             Ok(output) => {
                 if let Some(metrics) = output.spec_decode_metrics {
-                    spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics));
+                    spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(*metrics));
                 }
                 if prompt_tokens.is_none()
                     && let Some(info) = output.prompt_info.as_ref()
@@ -632,11 +632,11 @@ mod tests {
             }),
             Ok(GenerateOutput {
                 request_id: String::new(),
-                prompt_info: Some(GeneratePromptInfo {
+                prompt_info: Some(Box::new(GeneratePromptInfo {
                     prompt_token_ids: Arc::from([11_u32, 22_u32]),
                     prompt_logprobs: None,
                     prompt_token_id_logprobs: None,
-                }),
+                })),
                 token_ids: vec![33],
                 logprobs: None,
                 finish_reason: Some(FinishReason::stop_eos()),
@@ -716,10 +716,12 @@ mod tests {
     ) -> GenerateOutput {
         GenerateOutput {
             request_id: String::new(),
-            prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
-                prompt_token_ids: Arc::from(ids),
-                prompt_logprobs: None,
-                prompt_token_id_logprobs: None,
+            prompt_info: prompt_token_ids.map(|ids| {
+                Box::new(GeneratePromptInfo {
+                    prompt_token_ids: Arc::from(ids),
+                    prompt_logprobs: None,
+                    prompt_token_id_logprobs: None,
+                })
             }),
             token_ids,
             logprobs: None,
@@ -817,7 +819,7 @@ mod tests {
             ..Default::default()
         };
         let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos()));
-        output.spec_decode_metrics = Some(metrics.clone());
+        output.spec_decode_metrics = Some(Box::new(metrics.clone()));
 
         for include_usage in [false, true] {
             let chunks: Vec<_> = generate_chunk_stream(

rust/src/server/src/routes/tests.rs

@@ -120,8 +120,8 @@ fn request_output_with_logprobs_and_kv(
         new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct),
         finish_reason,
         stop_reason,
-        kv_transfer_params,
-        ec_transfer_params,
+        kv_transfer_params: kv_transfer_params.map(Box::new),
+        ec_transfer_params: ec_transfer_params.map(Box::new),
         ..Default::default()
     }
 }

rust/src/text/src/output/decoded.rs

@@ -47,10 +47,12 @@ pub struct Finished {
     pub usage: TokenUsage,
     pub finish_reason: FinishReason,
     /// Connector-specific KV transfer parameters for disaggregated serving.
-    pub kv_transfer_params: Option<serde_json::Value>,
+    /// Kept boxed as received from `vllm-llm`, so the chat events built from
+    /// this stay small without boxing again.
+    pub kv_transfer_params: Option<Box<serde_json::Value>>,
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
-    pub ec_transfer_params: Option<serde_json::Value>,
+    pub ec_transfer_params: Option<Box<serde_json::Value>>,
     /// Sampling support sets aligned with all generated token positions.
     pub sampling_mask: Option<SamplingMask>,
 }

rust/src/text/src/output/mod.rs

@@ -101,8 +101,10 @@ impl<T: TextOutputStream> T {
                             let mut collected = collected.unwrap();
                             collected.finish_reason = finished.finish_reason;
                             collected.usage = finished.usage;
-                            collected.kv_transfer_params = finished.kv_transfer_params;
-                            collected.ec_transfer_params = finished.ec_transfer_params;
+                            collected.kv_transfer_params =
+                                finished.kv_transfer_params.map(|value| *value);
+                            collected.ec_transfer_params =
+                                finished.ec_transfer_params.map(|value| *value);
                             collected.sampling_mask = finished.sampling_mask;
                             return Ok(collected);
                         }