@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 10e6a7f210

[Frontend] Expose per-request spec decode metrics in generate API (#43310)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Codex <codex@openai.com>

Hongxin Xu committed Sep 20, 2026 at 20:21 UTC 10e6a7f21094b76c65efb029b39b3f2227516418
12 files changed +516 -9
rust/src/engine-core-client/src/protocol/output.rs
+29 -3
@@ -130,10 +130,36 @@ pub struct EngineCoreOutput {
130 #[serde(default)]
131 pub new_sampling_mask: Option<MaybeWireSamplingMask>,
132 /// Per-request speculative-decoding acceptance metrics, set on the final
133 - /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here;
134 - /// the Rust frontend does not yet surface it in responses.
133 + /// output when `--per-request-spec-decode-metrics` is enabled.
134 #[serde(default)]
136 - pub spec_decode_metrics: Option<OpaqueValue>,
135 + pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
136 +}
137 +
138 +/// Raw per-sequence speculative-decoding accumulator.
139 +///
140 +/// Python models this as a plain `@dataclass`, so it is serialized by msgspec
141 +/// as a map (named fields). Derived values such as acceptance rates are
142 +/// computed by the frontend, not carried on the wire.
143 +///
144 +/// Original Python definition: `RequestSpecDecodeMetrics` in
145 +/// `vllm/v1/metrics/stats.py`.
146 +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
147 +pub struct RequestSpecDecodeMetrics {
148 + /// Configured `num_speculative_tokens`.
149 + #[serde(default)]
150 + pub num_spec_tokens: u64,
151 + /// Verify-step counts indexed by accepted draft-token count.
152 + #[serde(default)]
153 + pub histogram: Vec<u64>,
154 + /// Total proposed draft tokens.
155 + #[serde(default)]
156 + pub num_draft_tokens: u64,
157 + /// Accepted draft count per verify step; empty unless `detailed`.
158 + #[serde(default)]
159 + pub per_step_accepted: Vec<u64>,
160 + /// Proposed draft count per verify step; empty unless `detailed`.
161 + #[serde(default)]
162 + pub per_step_drafted: Vec<u64>,
163 }
164
165 impl EngineCoreOutput {
rust/src/llm/src/output.rs
+11 -1
@@ -11,7 +11,9 @@ use futures::stream::FusedStream;
11 use futures::{Stream, StreamExt as _, pin_mut};
12 use serde::{Deserialize, Serialize};
13 use vllm_engine_core_client::protocol::logprobs::Logprobs;
14 -use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
14 +use vllm_engine_core_client::protocol::output::{
15 + EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason,
16 +};
17 use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
18 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
19
@@ -47,6 +49,8 @@ pub struct CollectedGenerateOutput {
49 pub ec_transfer_params: Option<serde_json::Value>,
50 /// Sampling support sets aligned one-to-one with generated token positions.
51 pub sampling_mask: Option<SamplingMask>,
52 + /// Per-request speculative-decoding metrics from the terminal output.
53 + pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
54 }
55
56 /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
@@ -160,6 +164,8 @@ pub struct GenerateOutput {
164 pub ec_transfer_params: Option<serde_json::Value>,
165 /// Sampling support sets aligned one-to-one with `token_ids`.
166 pub sampling_mask: Option<SamplingMask>,
167 + /// Per-request speculative-decoding metrics, present on terminal outputs.
168 + pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
169 }
170
171 impl GenerateOutput {
@@ -208,6 +214,7 @@ impl GenerateOutput {
214 kv_transfer_params: None,
215 ec_transfer_params: None,
216 sampling_mask: None,
217 + spec_decode_metrics: None,
218 }
219 }
220 }
@@ -308,6 +315,7 @@ impl Stream for GenerateOutputStream {
315 kv_transfer_params: raw.kv_transfer_params,
316 ec_transfer_params: raw.ec_transfer_params,
317 sampling_mask,
318 + spec_decode_metrics: raw.spec_decode_metrics,
319 };
320
321 Poll::Ready(Some(Ok(output)))
@@ -396,6 +404,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
404 kv_transfer_params: None,
405 ec_transfer_params: None,
406 sampling_mask,
407 + spec_decode_metrics: None,
408 });
409 }
410
@@ -409,6 +418,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
418 };
419 collected.kv_transfer_params = output.kv_transfer_params;
420 collected.ec_transfer_params = output.ec_transfer_params;
421 + collected.spec_decode_metrics = output.spec_decode_metrics;
422 if let Some(mask) = collected.sampling_mask.as_ref()
423 && mask.rows.len() != collected.token_ids.len()
424 {
rust/src/llm/tests/generate.rs
+1
@@ -413,6 +413,7 @@ async fn collect_output_rejects_partial_sampling_mask() {
413 sampling_mask: Some(SamplingMask {
414 rows: vec![vec![1, 33, 99]],
415 }),
416 + spec_decode_metrics: None,
417 };
418
419 let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err();
rust/src/server/src/routes/inference/generate.rs
+140 -1
@@ -29,7 +29,8 @@ use vllm_llm::{
29 use self::convert::{ResponseOptions, prepare_generate_request};
30 use self::types::{
31 GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice,
32 - GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo,
32 + GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo,
33 + SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics,
34 };
35 pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams};
36 pub(crate) use self::validate::validate_request_compat;
@@ -160,10 +161,14 @@ async fn generate_chunk_stream(
161 let mut prompt_tokens = None;
162 let mut prompt_token_ids = None;
163 let mut usage = TokenUsage::default();
164 + let mut spec_decode_metrics = None;
165
166 while let Some(next) = stream.next().await {
167 match next {
168 Ok(output) => {
169 + if let Some(metrics) = output.spec_decode_metrics {
170 + spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics));
171 + }
172 if prompt_tokens.is_none()
173 && let Some(info) = output.prompt_info.as_ref()
174 {
@@ -223,6 +228,7 @@ async fn generate_chunk_stream(
228 .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
229 mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()),
230 prompt_token_ids,
231 + metrics: None,
232 })
233 .await;
234 }
@@ -243,6 +249,11 @@ async fn generate_chunk_stream(
249 usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
250 prompt_token_ids: None,
251 mm_placeholders: None,
252 + metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics {
253 + speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
254 + speculative_decoding,
255 + ),
256 + }),
257 })
258 .await;
259 }
@@ -318,6 +329,9 @@ fn collect_generate(
329 mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(),
330 kv_transfer_params: collected.kv_transfer_params,
331 ec_transfer_params: collected.ec_transfer_params,
332 + metrics: collected.spec_decode_metrics.map(|metrics| PerRequestMetrics {
333 + speculative_decoding: SpeculativeDecodingMetrics::from(metrics),
334 + }),
335 })
336 }
337
@@ -465,6 +479,7 @@ mod tests {
479
480 use futures::{TryStreamExt as _, stream};
481 use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};
482 + use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
483 use vllm_llm::GeneratePromptInfo;
484
485 use super::*;
@@ -482,6 +497,7 @@ mod tests {
497 kv_transfer_params: None,
498 ec_transfer_params: None,
499 sampling_mask: None,
500 + spec_decode_metrics: None,
501 }),
502 Ok(GenerateOutput {
503 request_id: String::new(),
@@ -496,6 +512,7 @@ mod tests {
512 kv_transfer_params: None,
513 ec_transfer_params: None,
514 sampling_mask: None,
515 + spec_decode_metrics: None,
516 }),
517 ]);
518
@@ -578,6 +595,7 @@ mod tests {
595 kv_transfer_params: None,
596 ec_transfer_params: None,
597 sampling_mask: None,
598 + spec_decode_metrics: None,
599 }
600 }
601
@@ -645,6 +663,80 @@ mod tests {
663 assert!(chunks[1].mm_placeholders.is_none());
664 }
665
666 + #[tokio::test]
667 + async fn generate_chunk_stream_returns_terminal_spec_decode_metrics() {
668 + let metrics = RequestSpecDecodeMetrics {
669 + num_spec_tokens: 3,
670 + histogram: vec![0, 1, 1, 0],
671 + num_draft_tokens: 6,
672 + ..Default::default()
673 + };
674 + let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos()));
675 + output.spec_decode_metrics = Some(metrics.clone());
676 +
677 + for include_usage in [false, true] {
678 + let chunks: Vec<_> = generate_chunk_stream(
679 + stream::iter(vec![Ok(output.clone())]),
680 + "raw-stream".to_string(),
681 + ApiServerOptions::default(),
682 + ResponseOptions {
683 + include_usage,
684 + ..Default::default()
685 + },
686 + None,
687 + )
688 + .try_collect()
689 + .await
690 + .expect("collect chunks");
691 + assert_eq!(chunks.len(), 1 + usize::from(include_usage));
692 + assert!(chunks[0].metrics.is_none());
693 + if include_usage {
694 + assert!(chunks[1].choices.is_empty());
695 + assert_eq!(
696 + chunks[1].metrics,
697 + Some(PerRequestMetrics {
698 + speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
699 + SpeculativeDecodingMetrics::from(metrics.clone())
700 + )
701 + })
702 + );
703 + let json = serde_json::to_value(&chunks[1]).expect("serialize chunk");
704 + let stats = &json["metrics"]["speculative_decoding"];
705 + assert_eq!(stats.as_object().expect("metrics object").len(), 7);
706 + assert_eq!(stats["num_draft_tokens"], 6);
707 + assert_eq!(stats["num_accepted_draft_tokens"], 3);
708 + assert_eq!(stats["mean_acceptance_length"], 2.5);
709 + assert_eq!(
710 + stats["acceptance_histogram"],
711 + serde_json::json!([0, 1, 1, 0])
712 + );
713 + assert!(stats.get("per_step_accepted").is_none());
714 + }
715 + }
716 + }
717 +
718 + #[test]
719 + fn spec_decode_metrics_zero_steps_and_detailed() {
720 + let empty = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
721 + num_spec_tokens: 2,
722 + histogram: vec![0, 0, 0],
723 + ..Default::default()
724 + });
725 + assert_eq!(empty.mean_acceptance_length, 1.0);
726 + assert_eq!(empty.draft_acceptance_rate, 0.0);
727 + assert_eq!(empty.per_step_accepted, None);
728 +
729 + let detailed = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
730 + num_spec_tokens: 2,
731 + histogram: vec![0, 1, 1],
732 + num_draft_tokens: 4,
733 + per_step_accepted: vec![1, 2],
734 + per_step_drafted: vec![2, 2],
735 + });
736 + assert_eq!(detailed.per_step_accepted, Some(vec![1, 2]));
737 + assert_eq!(detailed.per_step_drafted, Some(vec![2, 2]));
738 + }
739 +
740 #[test]
741 fn collect_generate_omits_prompt_metadata_by_default() {
742 let output = CollectedGenerateOutput {
@@ -662,6 +754,7 @@ mod tests {
754 ec_transfer_params: None,
755 prompt_token_ids: vec![10, 20],
756 sampling_mask: None,
757 + spec_decode_metrics: None,
758 };
759
760 let response = collect_generate(
@@ -683,6 +776,50 @@ mod tests {
776 assert!(response.mm_placeholders.is_none());
777 }
778
779 + #[test]
780 + fn collect_generate_returns_spec_decode_metrics() {
781 + let metrics = RequestSpecDecodeMetrics {
782 + num_spec_tokens: 3,
783 + histogram: vec![0, 1, 1, 0],
784 + num_draft_tokens: 6,
785 + ..Default::default()
786 + };
787 + let output = CollectedGenerateOutput {
788 + request_id: "raw-metrics".to_string(),
789 + prompt_logprobs: None,
790 + token_ids: vec![30],
791 + logprobs: None,
792 + finish_reason: FinishReason::stop_eos(),
793 + usage: TokenUsage::default(),
794 + kv_transfer_params: None,
795 + ec_transfer_params: None,
796 + prompt_token_ids: vec![10],
797 + sampling_mask: None,
798 + spec_decode_metrics: Some(metrics.clone()),
799 + };
800 +
801 + let response = collect_generate(
802 + output,
803 + "raw-metrics".to_string(),
804 + ApiServerOptions::default(),
805 + ResponseOptions::default(),
806 + None,
807 + )
808 + .expect("response");
809 +
810 + assert_eq!(
811 + response.metrics,
812 + Some(PerRequestMetrics {
813 + speculative_decoding: SpeculativeDecodingMetrics::from(metrics)
814 + })
815 + );
816 + let json = serde_json::to_value(&response).expect("serialize response");
817 + let stats = &json["metrics"]["speculative_decoding"];
818 + assert_eq!(stats.as_object().expect("metrics object").len(), 9);
819 + assert!(stats["per_step_accepted"].is_null());
820 + assert!(stats["per_step_drafted"].is_null());
821 + }
822 +
823 #[test]
824 fn extract_mm_placeholders_returns_none_without_features() {
825 assert!(extract_mm_placeholders(None).is_none());
@@ -749,6 +886,7 @@ mod tests {
886 ec_transfer_params: None,
887 prompt_token_ids: vec![10, 20],
888 sampling_mask: None,
889 + spec_decode_metrics: None,
890 };
891
892 let response = collect_generate(
@@ -782,6 +920,7 @@ mod tests {
920 kv_transfer_params: None,
921 ec_transfer_params: None,
922 sampling_mask: None,
923 + spec_decode_metrics: None,
924 prompt_token_ids,
925 };
926
rust/src/server/src/routes/inference/generate/types.rs
+87
@@ -7,6 +7,7 @@ use llm_multimodal::MediaContentPart;
7 use serde::{Deserialize, Serialize};
8 use serde_json::{Map, Value};
9 use validator::Validate;
10 +use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
11 use vllm_text::SamplingParams;
12
13 use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
@@ -84,6 +85,7 @@ pub(super) struct GenerateStreamResponse {
85 pub usage: Option<Usage>,
86 pub prompt_token_ids: Option<Vec<u32>>,
87 pub mm_placeholders: Option<MultiModalPlaceholders>,
88 + pub metrics: Option<PerRequestMetrics<StreamingSpeculativeDecodingMetrics>>,
89 }
90
91 /// Mirrors the Python vLLM `GenerateResponse` class.
@@ -96,6 +98,91 @@ pub(super) struct GenerateResponse {
98 pub mm_placeholders: Option<MultiModalPlaceholders>,
99 pub kv_transfer_params: Option<Value>,
100 pub ec_transfer_params: Option<Value>,
101 + pub metrics: Option<PerRequestMetrics<SpeculativeDecodingMetrics>>,
102 +}
103 +
104 +#[derive(Debug, Clone, PartialEq, Serialize)]
105 +pub(super) struct PerRequestMetrics<T> {
106 + pub speculative_decoding: T,
107 +}
108 +
109 +/// Mirrors the Python vLLM `SpeculativeDecodingMetrics` class.
110 +///
111 +/// Derived from the raw engine accumulator the same way as Python
112 +/// `RequestSpecDecodeMetrics.to_dict`.
113 +#[derive(Debug, Clone, PartialEq, Serialize)]
114 +pub(super) struct SpeculativeDecodingMetrics {
115 + pub mean_acceptance_length: f64,
116 + pub draft_acceptance_rate: f64,
117 + pub acceptance_histogram: Vec<u64>,
118 + pub num_spec_steps: u64,
119 + pub num_accepted_draft_tokens: u64,
120 + pub num_draft_tokens: u64,
121 + pub num_spec_tokens: u64,
122 + pub per_step_accepted: Option<Vec<u64>>,
123 + pub per_step_drafted: Option<Vec<u64>>,
124 +}
125 +
126 +impl From<RequestSpecDecodeMetrics> for SpeculativeDecodingMetrics {
127 + fn from(raw: RequestSpecDecodeMetrics) -> Self {
128 + let num_spec_steps: u64 = raw.histogram.iter().sum();
129 + let num_accepted_draft_tokens: u64 =
130 + (0u64..).zip(&raw.histogram).map(|(accepted, count)| accepted * count).sum();
131 + let ratio = |num: u64, den: u64| {
132 + if den == 0 {
133 + 0.0
134 + } else {
135 + num as f64 / den as f64
136 + }
137 + };
138 + let detailed = !raw.per_step_accepted.is_empty();
139 + Self {
140 + mean_acceptance_length: if num_spec_steps == 0 {
141 + 1.0
142 + } else {
143 + 1.0 + ratio(num_accepted_draft_tokens, num_spec_steps)
144 + },
145 + draft_acceptance_rate: ratio(num_accepted_draft_tokens, raw.num_draft_tokens),
146 + acceptance_histogram: raw.histogram,
147 + num_spec_steps,
148 + num_accepted_draft_tokens,
149 + num_draft_tokens: raw.num_draft_tokens,
150 + num_spec_tokens: raw.num_spec_tokens,
151 + per_step_accepted: detailed.then_some(raw.per_step_accepted),
152 + per_step_drafted: detailed.then_some(raw.per_step_drafted),
153 + }
154 + }
155 +}
156 +
157 +/// Streaming form omits detailed fields when summary metrics are requested.
158 +#[serde_with::skip_serializing_none]
159 +#[derive(Debug, Clone, PartialEq, Serialize)]
160 +pub(super) struct StreamingSpeculativeDecodingMetrics {
161 + pub mean_acceptance_length: f64,
162 + pub draft_acceptance_rate: f64,
163 + pub acceptance_histogram: Vec<u64>,
164 + pub num_spec_steps: u64,
165 + pub num_accepted_draft_tokens: u64,
166 + pub num_draft_tokens: u64,
167 + pub num_spec_tokens: u64,
168 + pub per_step_accepted: Option<Vec<u64>>,
169 + pub per_step_drafted: Option<Vec<u64>>,
170 +}
171 +
172 +impl From<SpeculativeDecodingMetrics> for StreamingSpeculativeDecodingMetrics {
173 + fn from(metrics: SpeculativeDecodingMetrics) -> Self {
174 + Self {
175 + mean_acceptance_length: metrics.mean_acceptance_length,
176 + draft_acceptance_rate: metrics.draft_acceptance_rate,
177 + acceptance_histogram: metrics.acceptance_histogram,
178 + num_spec_steps: metrics.num_spec_steps,
179 + num_accepted_draft_tokens: metrics.num_accepted_draft_tokens,
180 + num_draft_tokens: metrics.num_draft_tokens,
181 + num_spec_tokens: metrics.num_spec_tokens,
182 + per_step_accepted: metrics.per_step_accepted,
183 + per_step_drafted: metrics.per_step_drafted,
184 + }
185 + }
186 }
187
188 pub(super) type MultiModalPlaceholders = HashMap<String, Vec<PlaceholderRangeInfo>>;
tests/entrypoints/scale_out/derender/test_derender.py
+45
@@ -52,6 +52,7 @@ def _make_generate_response(
52 logprobs: dict | None = None,
53 prompt_logprobs: list | None = None,
54 kv_transfer_params: dict | None = None,
55 + metrics: dict | None = None,
56 ) -> dict:
57 choice: dict = {
58 "index": 0,
@@ -64,6 +65,7 @@ def _make_generate_response(
65 "choices": [choice],
66 "prompt_logprobs": prompt_logprobs,
67 "kv_transfer_params": kv_transfer_params,
68 + "metrics": metrics,
69 }
70
71
@@ -266,6 +268,49 @@ async def test_derender_chat_kv_transfer_params_passthrough(client):
268 assert response.json()["kv_transfer_params"] == kv
269
270
271 +@pytest.mark.asyncio
272 +async def test_derender_chat_metrics_passthrough(client):
273 + gen_req = await _render_chat(client)
274 + metrics = {
275 + "speculative_decoding": {
276 + "mean_acceptance_length": 2.0,
277 + "draft_acceptance_rate": 0.5,
278 + "acceptance_histogram": [0, 1],
279 + "num_spec_steps": 1,
280 + "num_accepted_draft_tokens": 1,
281 + "num_draft_tokens": 2,
282 + "num_spec_tokens": 1,
283 + }
284 + }
285 + response = await client.post(
286 + "/v1/chat/completions/derender",
287 + json={
288 + "model": MODEL_NAME,
289 + "generate_response": _make_generate_response(
290 + gen_req["token_ids"][:3], metrics=metrics
291 + ),
292 + },
293 + )
294 + assert response.status_code == 200
295 + actual = response.json()["metrics"]
296 + assert set(actual) == {
297 + "time_to_first_token_ms",
298 + "generation_time_ms",
299 + "queue_time_ms",
300 + "mean_itl_ms",
301 + "tokens_per_second",
302 + "speculative_decoding",
303 + }
304 + assert all(
305 + actual[name] is None for name in actual if name != "speculative_decoding"
306 + )
307 + assert actual["speculative_decoding"] == {
308 + **metrics["speculative_decoding"],
309 + "per_step_accepted": None,
310 + "per_step_drafted": None,
311 + }
312 +
313 +
314 @pytest.mark.asyncio
315 async def test_derender_chat_empty_token_ids(client):
316 """Empty token_ids list returns 400."""
tests/entrypoints/scale_out/derender/test_derender_stream.py
+24
@@ -27,6 +27,8 @@ from vllm.entrypoints.generate.base.protocol import (
27 DeltaFunctionCall,
28 DeltaMessage,
29 DeltaToolCall,
30 + PerRequestMetrics,
31 + SpeculativeDecodingMetrics,
32 )
33 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
34 DerenderStreamState,
@@ -164,6 +166,7 @@ def _make_stream_chunk(
166 finish_reason: str | None = None,
167 request_id: str = "test-req",
168 usage: dict | None = None,
169 + metrics: PerRequestMetrics | None = None,
170 ) -> GenerateStreamResponse:
171 """Build a GenerateStreamResponse SSE chunk."""
172 return GenerateStreamResponse(
@@ -176,6 +179,7 @@ def _make_stream_chunk(
179 )
180 ],
181 usage=UsageInfo(**usage) if usage else None,
182 + metrics=metrics,
183 )
184
185
@@ -512,6 +516,26 @@ class TestDerenderChatStream:
516 assert chunk1.choices[0].delta.role == "assistant"
517 assert chunk2.choices[0].delta.role is None
518
519 + @pytest.mark.asyncio
520 + async def test_metrics_passthrough(self, derenderer):
521 + metrics = PerRequestMetrics(
522 + speculative_decoding=SpeculativeDecodingMetrics(
523 + mean_acceptance_length=2.0,
524 + draft_acceptance_rate=0.5,
525 + acceptance_histogram=[0, 1],
526 + num_spec_steps=1,
527 + num_accepted_draft_tokens=1,
528 + num_draft_tokens=2,
529 + num_spec_tokens=1,
530 + )
531 + )
532 + chunk, _ = await derenderer.derender_chat_stream(
533 + model=MODEL_NAME,
534 + generate_chunk=_make_stream_chunk([], metrics=metrics),
535 + state=DerenderStreamState(),
536 + )
537 + assert chunk.metrics == metrics
538 +
539 @pytest.mark.asyncio
540 async def test_chunked_equals_oneshot(self, derenderer, tokenizer):
541 """Sum of streaming content deltas == one shot decode."""
tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+142
@@ -28,6 +28,7 @@ from vllm.renderers import renderer_from_config
28 from vllm.renderers.online_renderer import OnlineRenderer
29 from vllm.sampling_params import SamplingParams
30 from vllm.v1.engine.async_llm import AsyncLLM
31 +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics
32
33 MODEL_NAME = "openai-community/gpt2"
34 BASE_MODEL_PATHS = [
@@ -130,6 +131,7 @@ def _make_request_output(
131 logprobs: list[dict[int, Any] | None] | None = None,
132 num_cached_tokens: int | None = None,
133 index: int = 0,
134 + spec_decode_metrics: RequestSpecDecodeMetrics | None = None,
135 ) -> RequestOutput:
136 return RequestOutput(
137 request_id=request_id,
@@ -144,6 +146,7 @@ def _make_request_output(
146 cumulative_logprob=None,
147 logprobs=logprobs,
148 finish_reason=finish_reason,
149 + spec_decode_metrics=spec_decode_metrics,
150 )
151 ],
152 finished=finished,
@@ -321,6 +324,145 @@ async def test_stream_basic():
324 assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
325
326
327 +@pytest.mark.asyncio
328 +@pytest.mark.parametrize("stream", [False, True])
329 +async def test_serve_tokens_returns_spec_decode_metrics(stream: bool):
330 + engine = _mock_engine()
331 + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
332 + metrics.observe(num_draft_tokens=3, num_accepted=2)
333 +
334 + async def mock_generate(*args, **kwargs):
335 + yield _make_request_output(
336 + "req-1",
337 + token_ids=[10],
338 + finish_reason="stop",
339 + finished=True,
340 + spec_decode_metrics=metrics,
341 + )
342 +
343 + engine.generate = MagicMock(side_effect=mock_generate)
344 + serving = _build_serving_tokens(engine)
345 + request = GenerateRequest(
346 + token_ids=[1, 2, 3],
347 + sampling_params=SamplingParams(max_tokens=1),
348 + model=MODEL_NAME,
349 + stream=stream,
350 + stream_options=StreamOptions(include_usage=True) if stream else None,
351 + )
352 +
353 + response = await serving.serve_tokens(request)
354 + if stream:
355 + chunks = [chunk async for chunk in response]
356 + payload = next(
357 + chunk
358 + for chunk in _parse_sse_chunks(chunks)
359 + if isinstance(chunk, dict) and chunk.get("metrics")
360 + )["metrics"]["speculative_decoding"]
361 + else:
362 + assert isinstance(response, GenerateResponse)
363 + assert response.metrics is not None
364 + assert response.metrics.speculative_decoding is not None
365 + payload = response.metrics.speculative_decoding.model_dump()
366 +
367 + assert payload["num_draft_tokens"] == 3
368 + assert payload["num_accepted_draft_tokens"] == 2
369 + assert payload["num_spec_steps"] == 1
370 +
371 +
372 +@pytest.mark.asyncio
373 +@pytest.mark.parametrize("include_usage", [False, True])
374 +async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output(
375 + include_usage,
376 +):
377 + engine = _mock_engine()
378 + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
379 + metrics.observe(num_draft_tokens=3, num_accepted=2)
380 +
381 + async def mock_generate(*args, **kwargs):
382 + yield _make_request_output("req-1", token_ids=[10])
383 + yield _make_request_output(
384 + "req-1",
385 + token_ids=[],
386 + finish_reason="stop",
387 + finished=True,
388 + spec_decode_metrics=metrics,
389 + )
390 +
391 + engine.generate = MagicMock(side_effect=mock_generate)
392 + serving = _build_serving_tokens(engine)
393 + request = GenerateRequest(
394 + token_ids=[1, 2, 3],
395 + sampling_params=SamplingParams(max_tokens=1),
396 + model=MODEL_NAME,
397 + stream=True,
398 + stream_options=StreamOptions(include_usage=include_usage),
399 + )
400 +
401 + chunks = [chunk async for chunk in await serving.serve_tokens(request)]
402 + data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"]
403 +
404 + assert "metrics" not in data_chunks[0]
405 + assert len(data_chunks) == 1 + int(include_usage)
406 + if include_usage:
407 + assert data_chunks[1]["choices"] == []
408 + assert (
409 + data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3
410 + )
411 +
412 +
413 +@pytest.mark.asyncio
414 +@pytest.mark.parametrize("stream", [False, True])
415 +async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling(
416 + stream: bool,
417 +):
418 + engine = _mock_engine()
419 + metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
420 + metrics.observe(num_draft_tokens=3, num_accepted=2)
421 + result = _make_request_output(
422 + "req-1",
423 + token_ids=[10],
424 + finish_reason="stop",
425 + finished=True,
426 + spec_decode_metrics=metrics,
427 + )
428 + result.outputs.append(
429 + CompletionOutput(
430 + index=1,
431 + text="",
432 + token_ids=[20],
433 + cumulative_logprob=None,
434 + logprobs=None,
435 + finish_reason="stop",
436 + spec_decode_metrics=metrics,
437 + )
438 + )
439 +
440 + async def mock_generate(*args, **kwargs):
441 + yield result
442 +
443 + engine.generate = MagicMock(side_effect=mock_generate)
444 + serving = _build_serving_tokens(engine)
445 + request = GenerateRequest(
446 + token_ids=[1, 2, 3],
447 + sampling_params=SamplingParams(max_tokens=1, n=2),
448 + model=MODEL_NAME,
449 + stream=stream,
450 + stream_options=StreamOptions(include_usage=True) if stream else None,
451 + )
452 +
453 + response = await serving.serve_tokens(request)
454 + if stream:
455 + chunks = [chunk async for chunk in response]
456 + assert all(
457 + "metrics" not in chunk
458 + for chunk in _parse_sse_chunks(chunks)
459 + if isinstance(chunk, dict)
460 + )
461 + else:
462 + assert isinstance(response, GenerateResponse)
463 + assert response.metrics is None
464 +
465 +
466 @pytest.mark.asyncio
467 async def test_non_stream_returns_final_prompt_token_ids_when_requested():
468 engine = _mock_engine()
vllm/entrypoints/scale_out/derender/serving.py
+4
@@ -179,6 +179,7 @@ class ServingDerender(BaseServing):
179 usage=usage,
180 prompt_logprobs=gen.prompt_logprobs,
181 kv_transfer_params=gen.kv_transfer_params,
182 + metrics=gen.metrics,
183 )
184
185 async def derender_completion_response(
@@ -245,6 +246,9 @@ class ServingDerender(BaseServing):
246 choices=choices,
247 usage=usage,
248 kv_transfer_params=kv_params,
249 + # Metrics describe one prompt. Parallel samples are already
250 + # suppressed by /generate; multi-prompt responses cannot be merged.
251 + metrics=first.metrics if len(request.generate_responses) == 1 else None,
252 )
253
254 async def derender_chat_stream_response(
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+7 -1
@@ -11,7 +11,11 @@ from pydantic import (
11 )
12
13 from vllm.config import ModelConfig
14 -from vllm.entrypoints.generate.base.protocol import StreamOptions, validate_cache_salt
14 +from vllm.entrypoints.generate.base.protocol import (
15 + PerRequestMetrics,
16 + StreamOptions,
17 + validate_cache_salt,
18 +)
19 from vllm.entrypoints.openai.chat_completion.protocol import (
20 ChatCompletionLogProbs,
21 ChatCompletionRequest,
@@ -360,6 +364,7 @@ class GenerateStreamResponse(BaseModel):
364 usage: UsageInfo | None = Field(default=None)
365 prompt_token_ids: list[int] | None = None
366 mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
367 + metrics: PerRequestMetrics | None = None
368
369
370 class GenerateResponse(BaseModel):
@@ -378,6 +383,7 @@ class GenerateResponse(BaseModel):
383 prompt_logprobs: list[dict[int, Logprob] | None] | None = None
384 prompt_token_ids: list[int] | None = None
385 mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
386 + metrics: PerRequestMetrics | None = None
387
388 kv_transfer_params: dict[str, Any] | None = Field(
389 default=None,
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+23 -3
@@ -12,9 +12,13 @@ from fastapi import Request
12
13 from vllm.engine.protocol import EngineClient
14 from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker
15 -from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
15 +from vllm.entrypoints.generate.base.protocol import (
16 + PerRequestMetrics,
17 + RequestResponseMetadata,
18 +)
19 from vllm.entrypoints.generate.base.serving import (
20 GenerateBaseServing,
21 + build_spec_decoding_metrics,
22 clamp_prompt_logprobs,
23 )
24 from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -373,6 +377,11 @@ class ServingTokens(GenerateBaseServing):
377
378 request_metadata.final_usage_info = usage
379
380 + per_request_metrics = None
381 + if request.sampling_params.n == 1:
382 + spec_stats = build_spec_decoding_metrics(final_res)
383 + if spec_stats is not None:
384 + per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats)
385 response = GenerateResponse(
386 request_id=request_id,
387 created=created_time,
@@ -384,6 +393,7 @@ class ServingTokens(GenerateBaseServing):
393 final_res.prompt_token_ids if request.return_token_ids else None
394 ),
395 mm_placeholders=request._response_mm_placeholders,
396 + metrics=per_request_metrics,
397 kv_transfer_params=final_res.kv_transfer_params,
398 ec_transfer_params=final_res.ec_transfer_params,
399 )
@@ -423,6 +433,7 @@ class ServingTokens(GenerateBaseServing):
433 prompt_token_ids: list[int] | None = None
434 num_cached_tokens = None
435 sampling_params: SamplingParams = request.sampling_params
436 + last_res: RequestOutput | None = None
437
438 include_usage, include_continuous_usage = should_include_usage(
439 request.stream_options, False
@@ -430,6 +441,7 @@ class ServingTokens(GenerateBaseServing):
441
442 try:
443 async for res in result_generator:
444 + last_res = res
445 if first_iteration:
446 if res.prompt_token_ids is not None:
447 num_prompt_tokens = len(res.prompt_token_ids)
@@ -496,10 +508,10 @@ class ServingTokens(GenerateBaseServing):
508 total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
509 )
510
499 - # Omit absent prompt metadata (Rust skips None fields).
511 + # Omit fields that are absent from token-bearing chunks.
512 exclude = {
513 name
502 - for name in ("prompt_token_ids", "mm_placeholders")
514 + for name in ("prompt_token_ids", "mm_placeholders", "metrics")
515 if getattr(chunk, name) is None
516 }
517 yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n"
@@ -517,10 +529,18 @@ class ServingTokens(GenerateBaseServing):
529 )
530
531 if include_usage:
532 + per_request_metrics = None
533 + if sampling_params.n == 1:
534 + spec_stats = build_spec_decoding_metrics(last_res)
535 + if spec_stats is not None:
536 + per_request_metrics = PerRequestMetrics(
537 + speculative_decoding=spec_stats
538 + )
539 final_chunk = GenerateStreamResponse(
540 request_id=request_id,
541 choices=[],
542 usage=final_usage_info,
543 + metrics=per_request_metrics,
544 )
545 yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
546
vllm/renderers/online_derenderer.py
+3
@@ -398,6 +398,7 @@ class OnlineDerenderer:
398 model=model,
399 choices=stream_choices,
400 usage=usage,
401 + metrics=generate_chunk.metrics,
402 )
403 return chunk, updated_state
404
@@ -595,6 +596,7 @@ class OnlineDerenderer:
596 model=model,
597 choices=stream_choices,
598 usage=usage,
599 + metrics=generate_chunk.metrics,
600 )
601 return chunk, updated_state
602
@@ -741,6 +743,7 @@ class OnlineDerenderer:
743 model=model,
744 choices=stream_choices,
745 usage=usage,
746 + metrics=generate_chunk.metrics,
747 )
748 return chunk, updated_state
749