diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index e57b8bbb86..5fe25a9969 100644 --- a/rust/src/chat/src/event.rs +++ b/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, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 8e62f3e891..8c1dae37f4 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/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, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 5f5bbf32fa..80035db4cb 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -150,8 +150,8 @@ impl StructuredEventState { &mut self, usage: ChatTokenUsage, finish_reason: FinishReason, - kv_transfer_params: Option, - ec_transfer_params: Option, + kv_transfer_params: Option>, + ec_transfer_params: Option>, ) -> Result> { let mut events = Vec::new(); self.close_open_text_block(&mut events); diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 72b913078c..509628e8e2 100644 --- a/rust/src/chat/src/stream.rs +++ b/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, .. } => { diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c33062eeb3..9339c117a8 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -82,6 +82,10 @@ pub struct EngineCoreEvent { /// /// Original Python definition: /// +/// +/// 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, #[serde(default)] - pub pooling_output: Option, + pub pooling_output: Option>, #[serde(default)] pub finish_reason: Option, #[serde(default)] @@ -103,17 +107,17 @@ pub struct EngineCoreOutput { #[serde(default)] pub events: Option>, #[serde(default)] - pub kv_transfer_params: Option, + pub kv_transfer_params: Option>, #[serde(default)] - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, #[serde(default)] - pub trace_headers: Option, + pub trace_headers: Option>, /// 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, + pub prefill_stats: Option>, #[serde(default)] - pub routed_experts: Option, + pub routed_experts: Option>, /// 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, + pub spec_decode_metrics: Option>, /// 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, + pub prompt_token_id_logprobs: Option>, } /// Raw per-sequence speculative-decoding accumulator. diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4b6d033539..4dde96f2e2 100644 --- a/rust/src/llm/src/output.rs +++ b/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, + pub prompt_info: Option>, /// Newly produced token IDs for this step. pub token_ids: Vec, /// 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, + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned one-to-one with `token_ids`. pub sampling_mask: Option, /// Per-request speculative-decoding metrics, present on terminal outputs. - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option>, } 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, + pending_prompt_info: Option>, 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> + 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() { diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 98f9cadd1a..61cea12094 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/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() }, ); diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index c596c89ba3..813171555b 100644 --- a/rust/src/llm/tests/generate.rs +++ b/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], diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 4844aba120..8685357b18 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/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, diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 71b4286033..d82ceac083 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/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); diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index d495fea485..8e31bb84d2 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -406,8 +406,8 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result { 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( diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 0ed2ea4d45..a9c9831bbd 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/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() } } diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index d3dc422656..64df0492bf 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/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, + /// Kept boxed as received from `vllm-llm`, so the chat events built from + /// this stay small without boxing again. + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned with all generated token positions. pub sampling_mask: Option, } diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 1293ac6025..8e09f07618 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -101,8 +101,10 @@ impl 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); }