@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0983aef8da

[Frontend] Return prompt metadata from /inference/v1/generate (#53187)

aoshen02 committed Sep 17, 2026 at 06:29 UTC 0983aef8da9b25ae5be72a5567e571a8f3f82963
9 files changed +763 -40
rust/src/server/src/routes/inference/generate.rs
+268 -4
@@ -21,6 +21,7 @@ use thiserror_ext::AsReport as _;
21 use tracing::{error, info, trace};
22 use tracing_futures::Instrument as _;
23 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
24 +use vllm_engine_core_client::protocol::multimodal::MmFeatureSpec;
25 use vllm_llm::{
26 CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage,
27 };
@@ -28,7 +29,7 @@ use vllm_llm::{
29 use self::convert::{ResponseOptions, prepare_generate_request};
30 use self::types::{
31 GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice,
31 - GenerateStreamResponse,
32 + GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo,
33 };
34 pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams};
35 pub(crate) use self::validate::validate_request_compat;
@@ -64,6 +65,11 @@ pub async fn generate(
65 } else {
66 None
67 };
68 + let mm_placeholders = body
69 + .return_token_ids
70 + .unwrap_or(false)
71 + .then(|| extract_mm_placeholders(mm_features.as_deref()))
72 + .flatten();
73
74 let prepared =
75 match prepare_generate_request(body, &lora_resolution, request_context, mm_features) {
@@ -98,6 +104,7 @@ pub async fn generate(
104 prepared.request_id,
105 api_server_options,
106 prepared.options,
107 + mm_placeholders,
108 );
109 let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
110
@@ -120,6 +127,7 @@ pub async fn generate(
127 prepared.request_id,
128 api_server_options,
129 prepared.options,
130 + mm_placeholders,
131 ) {
132 Ok(response) => response,
133 Err(error) => return error.into_response(),
@@ -143,19 +151,26 @@ async fn generate_chunk_stream(
151 include_logprobs,
152 // Ignored: raw generate streaming has no prompt-logprobs wire shape.
153 include_prompt_logprobs: _,
154 + return_token_ids,
155 }: ResponseOptions,
156 + mut mm_placeholders: Option<MultiModalPlaceholders>,
157 mut y: TryYielder<GenerateStreamResponse, ApiError>,
158 ) -> Result<(), ApiError> {
159 pin_mut!(stream);
160 let mut prompt_tokens = None;
161 + let mut prompt_token_ids = None;
162 let mut usage = TokenUsage::default();
163
164 while let Some(next) = stream.next().await {
165 match next {
166 Ok(output) => {
156 - if prompt_tokens.is_none() {
157 - prompt_tokens =
158 - output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len());
167 + if prompt_tokens.is_none()
168 + && let Some(info) = output.prompt_info.as_ref()
169 + {
170 + prompt_tokens = Some(info.prompt_token_ids.len());
171 + if return_token_ids {
172 + prompt_token_ids = Some(info.prompt_token_ids.to_vec());
173 + }
174 }
175 usage.prompt_token_count = prompt_tokens.unwrap_or_default();
176 usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count);
@@ -195,6 +210,7 @@ async fn generate_chunk_stream(
210 None
211 };
212
213 + let prompt_token_ids = prompt_token_ids.take();
214 y.yield_ok(GenerateStreamResponse {
215 request_id: request_id.clone(),
216 choices: vec![GenerateResponseStreamChoice {
@@ -205,6 +221,8 @@ async fn generate_chunk_stream(
221 }],
222 usage: include_continuous_usage
223 .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
224 + mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()),
225 + prompt_token_ids,
226 })
227 .await;
228 }
@@ -223,6 +241,8 @@ async fn generate_chunk_stream(
241 request_id,
242 choices: Vec::new(),
243 usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
244 + prompt_token_ids: None,
245 + mm_placeholders: None,
246 })
247 .await;
248 }
@@ -244,7 +264,9 @@ fn collect_generate(
264 include_continuous_usage: _,
265 include_logprobs,
266 include_prompt_logprobs,
267 + return_token_ids,
268 }: ResponseOptions,
269 + mm_placeholders: Option<MultiModalPlaceholders>,
270 ) -> Result<GenerateResponse, ApiError> {
271 let logprobs = if include_logprobs {
272 let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
@@ -292,11 +314,34 @@ fn collect_generate(
314 token_ids: collected.token_ids,
315 }],
316 prompt_logprobs,
317 + prompt_token_ids: return_token_ids.then_some(collected.prompt_token_ids),
318 + mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(),
319 kv_transfer_params: collected.kv_transfer_params,
320 ec_transfer_params: collected.ec_transfer_params,
321 })
322 }
323
324 +fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option<MultiModalPlaceholders> {
325 + let features = features?;
326 + if features.is_empty() {
327 + return None;
328 + }
329 +
330 + let mut placeholders = MultiModalPlaceholders::new();
331 + for feature in features {
332 + placeholders.entry(feature.modality.as_str().to_string()).or_default().push(
333 + PlaceholderRangeInfo {
334 + offset: feature.mm_position.offset,
335 + length: feature.mm_position.length,
336 + },
337 + );
338 + }
339 + for ranges in placeholders.values_mut() {
340 + ranges.sort_by_key(|range| range.offset);
341 + }
342 + Some(placeholders)
343 +}
344 +
345 fn raw_logprobs_to_openai_chat(logprobs: &Logprobs) -> Result<ChatLogProbs, ApiError> {
346 let content = logprobs
347 .positions
@@ -419,6 +464,7 @@ mod tests {
464 use std::sync::Arc;
465
466 use futures::{TryStreamExt as _, stream};
467 + use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};
468 use vllm_llm::GeneratePromptInfo;
469
470 use super::*;
@@ -451,6 +497,13 @@ mod tests {
497 }),
498 ]);
499
500 + let mm_placeholders = MultiModalPlaceholders::from([(
501 + "image".to_string(),
502 + vec![PlaceholderRangeInfo {
503 + offset: 1,
504 + length: 4,
505 + }],
506 + )]);
507 let chunks: Vec<_> = generate_chunk_stream(
508 stream,
509 "raw-stream".to_string(),
@@ -461,14 +514,20 @@ mod tests {
514 ResponseOptions {
515 include_usage: true,
516 include_continuous_usage: true,
517 + return_token_ids: true,
518 ..Default::default()
519 },
520 + Some(mm_placeholders.clone()),
521 )
522 .try_collect()
523 .await
524 .expect("collect chunks");
525
526 assert_eq!(chunks.len(), 2);
527 + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22]));
528 + assert_eq!(chunks[0].mm_placeholders, Some(mm_placeholders));
529 + assert!(chunks[1].prompt_token_ids.is_none());
530 + assert!(chunks[1].mm_placeholders.is_none());
531 assert_eq!(
532 chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
533 2
@@ -499,6 +558,209 @@ mod tests {
558 );
559 }
560
561 + fn stream_output(
562 + prompt_token_ids: Option<&[u32]>,
563 + token_ids: Vec<u32>,
564 + finish_reason: Option<FinishReason>,
565 + ) -> GenerateOutput {
566 + GenerateOutput {
567 + request_id: String::new(),
568 + prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
569 + prompt_token_ids: Arc::from(ids),
570 + prompt_logprobs: None,
571 + }),
572 + token_ids,
573 + logprobs: None,
574 + finish_reason,
575 + cached_token_count: 0,
576 + kv_transfer_params: None,
577 + ec_transfer_params: None,
578 + }
579 + }
580 +
581 + async fn collect_chunks(
582 + outputs: Vec<GenerateOutput>,
583 + return_token_ids: bool,
584 + mm_placeholders: Option<MultiModalPlaceholders>,
585 + ) -> Vec<GenerateStreamResponse> {
586 + generate_chunk_stream(
587 + stream::iter(outputs.into_iter().map(Ok)),
588 + "raw-stream".to_string(),
589 + ApiServerOptions::default(),
590 + ResponseOptions {
591 + return_token_ids,
592 + ..Default::default()
593 + },
594 + mm_placeholders,
595 + )
596 + .try_collect()
597 + .await
598 + .expect("collect chunks")
599 + }
600 +
601 + #[tokio::test]
602 + async fn generate_chunk_stream_omits_prompt_metadata_by_default() {
603 + let chunks = collect_chunks(
604 + vec![
605 + stream_output(Some(&[11, 22]), vec![33], None),
606 + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())),
607 + ],
608 + false,
609 + Some(MultiModalPlaceholders::from([(
610 + "image".to_string(),
611 + vec![PlaceholderRangeInfo {
612 + offset: 1,
613 + length: 4,
614 + }],
615 + )])),
616 + )
617 + .await;
618 +
619 + assert_eq!(chunks.len(), 2);
620 + assert!(chunks.iter().all(|c| c.prompt_token_ids.is_none()));
621 + assert!(chunks.iter().all(|c| c.mm_placeholders.is_none()));
622 + }
623 +
624 + #[tokio::test]
625 + async fn generate_chunk_stream_sends_prompt_metadata_on_first_emitted_chunk() {
626 + // An empty first output is skipped and must not consume the metadata.
627 + let chunks = collect_chunks(
628 + vec![
629 + stream_output(Some(&[11, 22]), Vec::new(), None),
630 + stream_output(Some(&[11, 22]), vec![33], None),
631 + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())),
632 + ],
633 + true,
634 + None,
635 + )
636 + .await;
637 +
638 + assert_eq!(chunks.len(), 2);
639 + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22]));
640 + assert!(chunks[0].mm_placeholders.is_none());
641 + assert!(chunks[1].prompt_token_ids.is_none());
642 + assert!(chunks[1].mm_placeholders.is_none());
643 + }
644 +
645 + #[test]
646 + fn collect_generate_omits_prompt_metadata_by_default() {
647 + let output = CollectedGenerateOutput {
648 + request_id: "raw-1".to_string(),
649 + prompt_logprobs: None,
650 + token_ids: vec![30],
651 + logprobs: None,
652 + finish_reason: FinishReason::stop_eos(),
653 + usage: TokenUsage {
654 + prompt_token_count: 2,
655 + output_token_count: 1,
656 + cached_token_count: 0,
657 + },
658 + kv_transfer_params: None,
659 + ec_transfer_params: None,
660 + prompt_token_ids: vec![10, 20],
661 + };
662 +
663 + let response = collect_generate(
664 + output,
665 + "raw-1".to_string(),
666 + ApiServerOptions::default(),
667 + ResponseOptions::default(),
668 + Some(MultiModalPlaceholders::from([(
669 + "image".to_string(),
670 + vec![PlaceholderRangeInfo {
671 + offset: 0,
672 + length: 1,
673 + }],
674 + )])),
675 + )
676 + .expect("response");
677 +
678 + assert!(response.prompt_token_ids.is_none());
679 + assert!(response.mm_placeholders.is_none());
680 + }
681 +
682 + #[test]
683 + fn extract_mm_placeholders_returns_none_without_features() {
684 + assert!(extract_mm_placeholders(None).is_none());
685 + assert!(extract_mm_placeholders(Some(&[])).is_none());
686 + }
687 +
688 + #[test]
689 + fn extract_mm_placeholders_groups_ranges_in_prompt_order() {
690 + let feature = |modality: MmModality, offset: usize, length: usize| MmFeatureSpec {
691 + data: None,
692 + modality,
693 + identifier: format!("{}-{offset}", modality.as_str()),
694 + mm_position: PlaceholderRange {
695 + offset,
696 + length,
697 + is_embed: None,
698 + },
699 + mm_hash: None,
700 + };
701 + let features = vec![
702 + feature(MmModality::Image, 8, 2),
703 + feature(MmModality::Audio, 4, 3),
704 + feature(MmModality::Image, 1, 5),
705 + ];
706 +
707 + let placeholders = extract_mm_placeholders(Some(&features)).expect("metadata");
708 +
709 + assert_eq!(
710 + placeholders["image"],
711 + vec![
712 + PlaceholderRangeInfo {
713 + offset: 1,
714 + length: 5,
715 + },
716 + PlaceholderRangeInfo {
717 + offset: 8,
718 + length: 2,
719 + },
720 + ]
721 + );
722 + assert_eq!(
723 + placeholders["audio"],
724 + vec![PlaceholderRangeInfo {
725 + offset: 4,
726 + length: 3,
727 + }]
728 + );
729 + }
730 +
731 + #[test]
732 + fn collect_generate_returns_prompt_token_ids_when_requested() {
733 + let output = CollectedGenerateOutput {
734 + request_id: "raw-1".to_string(),
735 + prompt_logprobs: None,
736 + token_ids: vec![30],
737 + logprobs: None,
738 + finish_reason: FinishReason::stop_eos(),
739 + usage: TokenUsage {
740 + prompt_token_count: 2,
741 + output_token_count: 1,
742 + cached_token_count: 0,
743 + },
744 + kv_transfer_params: None,
745 + ec_transfer_params: None,
746 + prompt_token_ids: vec![10, 20],
747 + };
748 +
749 + let response = collect_generate(
750 + output,
751 + "raw-1".to_string(),
752 + ApiServerOptions::default(),
753 + ResponseOptions {
754 + return_token_ids: true,
755 + ..Default::default()
756 + },
757 + None,
758 + )
759 + .expect("response");
760 +
761 + assert_eq!(response.prompt_token_ids, Some(vec![10, 20]));
762 + }
763 +
764 #[test]
765 fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
766 let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
@@ -525,6 +787,7 @@ mod tests {
787 include_prompt_logprobs: true,
788 ..Default::default()
789 },
790 + None,
791 )
792 .expect("single-token prompt without payload maps to [None]");
793 let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
@@ -539,6 +802,7 @@ mod tests {
802 include_prompt_logprobs: true,
803 ..Default::default()
804 },
805 + None,
806 )
807 .expect_err("multi-token prompt without payload is an engine failure");
808 }
rust/src/server/src/routes/inference/generate/convert.rs
+4
@@ -30,6 +30,8 @@ pub(super) struct ResponseOptions {
30 pub include_logprobs: bool,
31 /// Whether the caller requested top-level prompt logprobs.
32 pub include_prompt_logprobs: bool,
33 + /// Whether the caller requested final prompt token metadata.
34 + pub return_token_ids: bool,
35 }
36
37 /// Validate and lower one raw generate request into the internal
@@ -56,6 +58,7 @@ pub(super) fn prepare_generate_request(
58 .unwrap_or(false);
59 let include_logprobs = request.sampling_params.inner.logprobs.is_some();
60 let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some();
61 + let return_token_ids = request.return_token_ids.unwrap_or(false);
62 let mut sampling_params = request.sampling_params.inner;
63 sampling_params.vllm_xargs = merge_kv_transfer_params(
64 sampling_params.vllm_xargs,
@@ -93,6 +96,7 @@ pub(super) fn prepare_generate_request(
96 include_continuous_usage,
97 include_logprobs,
98 include_prompt_logprobs,
99 + return_token_ids,
100 },
101 })
102 }
rust/src/server/src/routes/inference/generate/types.rs
+13
@@ -46,6 +46,7 @@ pub struct GenerateRequest {
46 pub ec_transfer_params: Option<HashMap<String, Value>>,
47 /// Raw multimodal input; server resolves media. Mutually exclusive with `features`.
48 pub content_parts: Option<Vec<MediaContentPart>>,
49 + pub return_token_ids: Option<bool>,
50 #[serde(flatten)]
51 pub other: Map<String, Value>,
52 }
@@ -81,6 +82,8 @@ pub(super) struct GenerateStreamResponse {
82 pub request_id: String,
83 pub choices: Vec<GenerateResponseStreamChoice>,
84 pub usage: Option<Usage>,
85 + pub prompt_token_ids: Option<Vec<u32>>,
86 + pub mm_placeholders: Option<MultiModalPlaceholders>,
87 }
88
89 /// Mirrors the Python vLLM `GenerateResponse` class.
@@ -89,10 +92,20 @@ pub(super) struct GenerateResponse {
92 pub request_id: String,
93 pub choices: Vec<GenerateResponseChoice>,
94 pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
95 + pub prompt_token_ids: Option<Vec<u32>>,
96 + pub mm_placeholders: Option<MultiModalPlaceholders>,
97 pub kv_transfer_params: Option<Value>,
98 pub ec_transfer_params: Option<Value>,
99 }
100
101 +pub(super) type MultiModalPlaceholders = HashMap<String, Vec<PlaceholderRangeInfo>>;
102 +
103 +#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
104 +pub(super) struct PlaceholderRangeInfo {
105 + pub offset: usize,
106 + pub length: usize,
107 +}
108 +
109 /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.
110 #[derive(Debug, Clone, Serialize)]
111 pub(super) struct GenerateLogprob {
rust/src/server/src/routes/render.rs
+1
@@ -105,6 +105,7 @@ fn lower_render_request(
105 kv_transfer_params: None,
106 ec_transfer_params: None,
107 content_parts: None,
108 + return_token_ids: None,
109 other: Default::default(),
110 };
111 validate_generate_request(&request, &state.served_model_names)?;
tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+357
@@ -12,6 +12,9 @@ from vllm.config.multimodal import MultiModalConfig
12 from vllm.entrypoints.generate.base.protocol import StreamOptions
13 from vllm.entrypoints.openai.models.protocol import BaseModelPath
14 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
15 +from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
16 + placeholder_ranges_from_engine_input,
17 +)
18 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
19 GenerateRequest,
20 GenerateResponse,
@@ -19,6 +22,7 @@ from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
22 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
23 from vllm.exceptions import GenerationError
24 from vllm.logprobs import Logprob
25 +from vllm.multimodal.inputs import PlaceholderRange
26 from vllm.outputs import CompletionOutput, RequestOutput
27 from vllm.renderers import renderer_from_config
28 from vllm.renderers.online_renderer import OnlineRenderer
@@ -176,6 +180,29 @@ def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
180 return parsed
181
182
183 +def test_placeholder_ranges_from_engine_input():
184 + engine_input: Any = {
185 + "type": "multimodal",
186 + "mm_placeholders": {
187 + "image": [
188 + PlaceholderRange(offset=1, length=5),
189 + PlaceholderRange(offset=8, length=2),
190 + ]
191 + },
192 + }
193 +
194 + placeholders = placeholder_ranges_from_engine_input(engine_input)
195 +
196 + assert placeholders is not None
197 + assert [item.model_dump() for item in placeholders["image"]] == [
198 + {"offset": 1, "length": 5},
199 + {"offset": 8, "length": 2},
200 + ]
201 +
202 + text_only: Any = {"type": "token", "prompt_token_ids": [1, 2, 3]}
203 + assert placeholder_ranges_from_engine_input(text_only) is None
204 +
205 +
206 @pytest.mark.asyncio
207 async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution():
208 engine = _mock_engine()
@@ -294,6 +321,336 @@ async def test_stream_basic():
321 assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
322
323
324 +@pytest.mark.asyncio
325 +async def test_non_stream_returns_final_prompt_token_ids_when_requested():
326 + engine = _mock_engine()
327 + final_prompt_token_ids = [1, 2, 2, 2, 3]
328 +
329 + async def mock_generate(*args, **kwargs):
330 + yield _make_request_output(
331 + "req-1",
332 + token_ids=[10],
333 + finish_reason="stop",
334 + finished=True,
335 + prompt_token_ids=final_prompt_token_ids,
336 + )
337 +
338 + engine.generate = MagicMock(side_effect=mock_generate)
339 + serving = _build_serving_tokens(engine)
340 + request = GenerateRequest(
341 + token_ids=[1, 2, 3],
342 + sampling_params=SamplingParams(max_tokens=1),
343 + model=MODEL_NAME,
344 + return_token_ids=True,
345 + )
346 +
347 + response = await serving.serve_tokens(request)
348 +
349 + assert isinstance(response, GenerateResponse)
350 + assert response.prompt_token_ids == final_prompt_token_ids
351 +
352 +
353 +@pytest.mark.asyncio
354 +async def test_stream_returns_prompt_token_ids_only_once_when_requested():
355 + engine = _mock_engine()
356 + final_prompt_token_ids = [1, 2, 2, 2, 3]
357 +
358 + async def mock_generate(*args, **kwargs):
359 + yield _make_request_output(
360 + "req-1", token_ids=[10], prompt_token_ids=final_prompt_token_ids
361 + )
362 + yield _make_request_output(
363 + "req-1",
364 + token_ids=[20],
365 + finish_reason="stop",
366 + finished=True,
367 + prompt_token_ids=final_prompt_token_ids,
368 + )
369 +
370 + engine.generate = MagicMock(side_effect=mock_generate)
371 + serving = _build_serving_tokens(engine)
372 + request = GenerateRequest(
373 + token_ids=[1, 2, 3],
374 + sampling_params=SamplingParams(max_tokens=2),
375 + model=MODEL_NAME,
376 + stream=True,
377 + return_token_ids=True,
378 + )
379 +
380 + response = await serving.serve_tokens(request)
381 + parsed = _parse_sse_chunks([chunk async for chunk in response])
382 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
383 +
384 + assert data_chunks[0]["prompt_token_ids"] == final_prompt_token_ids
385 + assert "prompt_token_ids" not in data_chunks[1]
386 +
387 +
388 +MM_ENGINE_INPUT: Any = {
389 + "type": "multimodal",
390 + "prompt_token_ids": [1, 2, 2, 2, 3],
391 + "mm_placeholders": {"image": [PlaceholderRange(offset=1, length=3)]},
392 +}
393 +MM_PLACEHOLDERS = {"image": [{"offset": 1, "length": 3}]}
394 +
395 +
396 +def _build_mm_serving_tokens(engine: AsyncLLM) -> ServingTokens:
397 + serving = _build_serving_tokens(engine)
398 +
399 + async def _fake_preprocess(*args, **kwargs):
400 + return [MM_ENGINE_INPUT]
401 +
402 + serving.online_renderer.preprocess_completion = AsyncMock(
403 + side_effect=_fake_preprocess
404 + )
405 + return serving
406 +
407 +
408 +@pytest.mark.asyncio
409 +@pytest.mark.parametrize("return_token_ids", [False, True])
410 +async def test_non_stream_prompt_metadata_requires_return_token_ids(
411 + return_token_ids: bool,
412 +):
413 + engine = _mock_engine()
414 +
415 + async def mock_generate(*args, **kwargs):
416 + yield _make_request_output(
417 + "req-1", token_ids=[10], finish_reason="stop", finished=True
418 + )
419 +
420 + engine.generate = MagicMock(side_effect=mock_generate)
421 + serving = _build_mm_serving_tokens(engine)
422 + request = GenerateRequest(
423 + token_ids=[1, 2, 3],
424 + sampling_params=SamplingParams(max_tokens=1),
425 + model=MODEL_NAME,
426 + return_token_ids=return_token_ids,
427 + )
428 +
429 + response = await serving.serve_tokens(request)
430 +
431 + assert isinstance(response, GenerateResponse)
432 + if return_token_ids:
433 + assert response.prompt_token_ids == [1, 2, 3]
434 + assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS
435 + else:
436 + assert response.prompt_token_ids is None
437 + assert response.mm_placeholders is None
438 + assert request._response_mm_placeholders is None
439 +
440 +
441 +@pytest.mark.asyncio
442 +async def test_non_stream_text_only_prompt_has_no_mm_placeholders():
443 + engine = _mock_engine()
444 +
445 + async def mock_generate(*args, **kwargs):
446 + yield _make_request_output(
447 + "req-1", token_ids=[10], finish_reason="stop", finished=True
448 + )
449 +
450 + engine.generate = MagicMock(side_effect=mock_generate)
451 + serving = _build_serving_tokens(engine)
452 + request = GenerateRequest(
453 + token_ids=[1, 2, 3],
454 + sampling_params=SamplingParams(max_tokens=1),
455 + model=MODEL_NAME,
456 + return_token_ids=True,
457 + )
458 +
459 + response = await serving.serve_tokens(request)
460 +
461 + assert isinstance(response, GenerateResponse)
462 + assert response.prompt_token_ids == [1, 2, 3]
463 + assert response.mm_placeholders is None
464 +
465 +
466 +@pytest.mark.asyncio
467 +async def test_non_stream_encoder_decoder_omits_mm_placeholders():
468 + engine = _mock_engine()
469 + engine.model_config.is_encoder_decoder = True
470 +
471 + async def mock_generate(*args, **kwargs):
472 + yield _make_request_output(
473 + "req-1", token_ids=[10], finish_reason="stop", finished=True
474 + )
475 +
476 + engine.generate = MagicMock(side_effect=mock_generate)
477 + serving = _build_mm_serving_tokens(engine)
478 + request = GenerateRequest(
479 + token_ids=[1, 2, 3],
480 + sampling_params=SamplingParams(max_tokens=1),
481 + model=MODEL_NAME,
482 + return_token_ids=True,
483 + )
484 +
485 + response = await serving.serve_tokens(request)
486 +
487 + assert isinstance(response, GenerateResponse)
488 + assert response.prompt_token_ids == [1, 2, 3]
489 + assert response.mm_placeholders is None
490 +
491 +
492 +@pytest.mark.asyncio
493 +async def test_stream_prompt_metadata_lands_on_first_emitted_chunk():
494 + """An empty first output is skipped; metadata goes on the first real chunk,
495 + once, and never on the sibling chunk of the same iteration (n=2)."""
496 + engine = _mock_engine()
497 +
498 + def _two_outputs(tokens0: list[int], tokens1: list[int], **kwargs):
499 + res = _make_request_output("req-1", token_ids=tokens0, **kwargs)
500 + res.outputs.append(
501 + CompletionOutput(
502 + index=1,
503 + text="",
504 + token_ids=tokens1,
505 + cumulative_logprob=None,
506 + logprobs=None,
507 + finish_reason=kwargs.get("finish_reason"),
508 + )
509 + )
510 + return res
511 +
512 + async def mock_generate(*args, **kwargs):
513 + yield _two_outputs([], [])
514 + yield _two_outputs([10], [11])
515 + yield _two_outputs([20], [21], finish_reason="stop", finished=True)
516 +
517 + engine.generate = MagicMock(side_effect=mock_generate)
518 + serving = _build_mm_serving_tokens(engine)
519 + request = GenerateRequest(
520 + token_ids=[1, 2, 3],
521 + sampling_params=SamplingParams(max_tokens=2, n=2),
522 + model=MODEL_NAME,
523 + stream=True,
524 + return_token_ids=True,
525 + )
526 +
527 + response = await serving.serve_tokens(request)
528 + parsed = _parse_sse_chunks([chunk async for chunk in response])
529 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
530 +
531 + assert [c["choices"][0]["token_ids"] for c in data_chunks] == [
532 + [10],
533 + [11],
534 + [20],
535 + [21],
536 + ]
537 + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
538 + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS
539 + for chunk in data_chunks[1:]:
540 + assert "prompt_token_ids" not in chunk
541 + assert "mm_placeholders" not in chunk
542 +
543 +
544 +@pytest.mark.asyncio
545 +async def test_stream_omits_prompt_metadata_by_default():
546 + engine = _mock_engine()
547 +
548 + async def mock_generate(*args, **kwargs):
549 + yield _make_request_output(
550 + "req-1", token_ids=[10], finish_reason="stop", finished=True
551 + )
552 +
553 + engine.generate = MagicMock(side_effect=mock_generate)
554 + serving = _build_mm_serving_tokens(engine)
555 + request = GenerateRequest(
556 + token_ids=[1, 2, 3],
557 + sampling_params=SamplingParams(max_tokens=1),
558 + model=MODEL_NAME,
559 + stream=True,
560 + )
561 +
562 + response = await serving.serve_tokens(request)
563 + parsed = _parse_sse_chunks([chunk async for chunk in response])
564 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
565 +
566 + assert data_chunks
567 + for chunk in data_chunks:
568 + assert "prompt_token_ids" not in chunk
569 + assert "mm_placeholders" not in chunk
570 +
571 +
572 +@pytest.mark.asyncio
573 +async def test_stream_text_only_prompt_omits_mm_placeholders_key():
574 + engine = _mock_engine()
575 +
576 + async def mock_generate(*args, **kwargs):
577 + yield _make_request_output(
578 + "req-1", token_ids=[10], finish_reason="stop", finished=True
579 + )
580 +
581 + engine.generate = MagicMock(side_effect=mock_generate)
582 + serving = _build_serving_tokens(engine)
583 + request = GenerateRequest(
584 + token_ids=[1, 2, 3],
585 + sampling_params=SamplingParams(max_tokens=1),
586 + model=MODEL_NAME,
587 + stream=True,
588 + return_token_ids=True,
589 + )
590 +
591 + response = await serving.serve_tokens(request)
592 + parsed = _parse_sse_chunks([chunk async for chunk in response])
593 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
594 +
595 + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
596 + assert "mm_placeholders" not in data_chunks[0]
597 +
598 +
599 +@pytest.mark.asyncio
600 +async def test_stream_zero_token_completion_still_delivers_prompt_metadata():
601 + engine = _mock_engine()
602 +
603 + async def mock_generate(*args, **kwargs):
604 + yield _make_request_output(
605 + "req-1", token_ids=[], finish_reason="stop", finished=True
606 + )
607 +
608 + engine.generate = MagicMock(side_effect=mock_generate)
609 + serving = _build_mm_serving_tokens(engine)
610 + request = GenerateRequest(
611 + token_ids=[1, 2, 3],
612 + sampling_params=SamplingParams(max_tokens=1),
613 + model=MODEL_NAME,
614 + stream=True,
615 + return_token_ids=True,
616 + )
617 +
618 + response = await serving.serve_tokens(request)
619 + parsed = _parse_sse_chunks([chunk async for chunk in response])
620 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
621 +
622 + assert len(data_chunks) == 1
623 + assert data_chunks[0]["choices"][0]["token_ids"] == []
624 + assert data_chunks[0]["choices"][0]["finish_reason"] == "stop"
625 + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
626 + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS
627 +
628 +
629 +@pytest.mark.asyncio
630 +async def test_stream_zero_token_completion_emits_no_chunk_by_default():
631 + engine = _mock_engine()
632 +
633 + async def mock_generate(*args, **kwargs):
634 + yield _make_request_output(
635 + "req-1", token_ids=[], finish_reason="stop", finished=True
636 + )
637 +
638 + engine.generate = MagicMock(side_effect=mock_generate)
639 + serving = _build_serving_tokens(engine)
640 + request = GenerateRequest(
641 + token_ids=[1, 2, 3],
642 + sampling_params=SamplingParams(max_tokens=1),
643 + model=MODEL_NAME,
644 + stream=True,
645 + )
646 +
647 + response = await serving.serve_tokens(request)
648 + parsed = _parse_sse_chunks([chunk async for chunk in response])
649 +
650 + assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")]
651 + assert parsed[-1] == "[DONE]"
652 +
653 +
654 @pytest.mark.asyncio
655 async def test_stream_error_mid_generation():
656 """finish_reason='error' mid-stream yields error chunk then [DONE]."""
tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py
+49 -25
@@ -154,11 +154,28 @@ async def test_render_to_generate_roundtrip(client, test_image):
154 )
155
156
157 -@pytest.mark.asyncio
158 -async def test_content_parts_generates_tokens(client, test_image):
159 - """content_parts with raw media should produce output tokens."""
160 - data_url = encode_image_url(test_image, format="PNG")
157 +def _collapse_placeholders(
158 + token_ids: list[int], mm_placeholders: dict[str, list[dict[str, int]]]
159 +) -> list[int]:
160 + """Undo placeholder expansion: each expanded range becomes one token.
161 +
162 + ``/generate`` with ``content_parts`` expects unexpanded placeholders and
163 + expands them itself, so the roundtrip must start from a collapsed prompt.
164 + """
165 + ranges = sorted(
166 + (r["offset"], r["length"]) for rs in mm_placeholders.values() for r in rs
167 + )
168 + collapsed: list[int] = []
169 + pos = 0
170 + for offset, length in ranges:
171 + collapsed.extend(token_ids[pos:offset])
172 + collapsed.append(token_ids[offset])
173 + pos = offset + length
174 + collapsed.extend(token_ids[pos:])
175 + return collapsed
176
177 +
178 +async def _render_image_prompt(client, data_url: str) -> dict:
179 render_resp = await client.post(
180 RENDER_ENDPOINT,
181 json={
@@ -175,14 +192,26 @@ async def test_content_parts_generates_tokens(client, test_image):
192 },
193 )
194 render_resp.raise_for_status()
178 - token_ids = render_resp.json()["token_ids"]
195 + return render_resp.json()
196 +
197 +
198 +@pytest.mark.asyncio
199 +async def test_content_parts_generates_tokens(client, test_image):
200 + """Unexpanded prompt + content_parts must expand exactly like /render."""
201 + data_url = encode_image_url(test_image, format="PNG")
202 + render_data = await _render_image_prompt(client, data_url)
203 + expanded = render_data["token_ids"]
204 + mm_placeholders = render_data["features"]["mm_placeholders"]
205 + unexpanded = _collapse_placeholders(expanded, mm_placeholders)
206 + assert len(unexpanded) < len(expanded)
207
208 gen_resp = await client.post(
209 GEN_ENDPOINT,
210 json={
183 - "token_ids": token_ids,
211 + "token_ids": unexpanded,
212 "content_parts": [{"type": "image_url", "url": data_url}],
213 "sampling_params": {"max_tokens": 10, "temperature": 0.0},
214 + "return_token_ids": True,
215 },
216 )
217 gen_resp.raise_for_status()
@@ -192,6 +221,8 @@ async def test_content_parts_generates_tokens(client, test_image):
221 choice = gen_data["choices"][0]
222 assert "token_ids" in choice
223 assert len(choice["token_ids"]) > 0
224 + assert gen_data["prompt_token_ids"] == expanded
225 + assert gen_data["mm_placeholders"] == mm_placeholders
226
227
228 @pytest.mark.asyncio
@@ -200,33 +231,20 @@ async def test_content_parts_streaming(client, test_image):
231 import json as json_mod
232
233 data_url = encode_image_url(test_image, format="PNG")
203 -
204 - render_resp = await client.post(
205 - RENDER_ENDPOINT,
206 - json={
207 - "model": MODEL_NAME,
208 - "messages": [
209 - {
210 - "role": "user",
211 - "content": [
212 - {"type": "image_url", "image_url": {"url": data_url}},
213 - {"type": "text", "text": "Describe this."},
214 - ],
215 - }
216 - ],
217 - },
218 - )
219 - render_resp.raise_for_status()
220 - token_ids = render_resp.json()["token_ids"]
234 + render_data = await _render_image_prompt(client, data_url)
235 + expanded = render_data["token_ids"]
236 + mm_placeholders = render_data["features"]["mm_placeholders"]
237 + unexpanded = _collapse_placeholders(expanded, mm_placeholders)
238
239 async with client.stream(
240 "POST",
241 GEN_ENDPOINT,
242 json={
226 - "token_ids": token_ids,
243 + "token_ids": unexpanded,
244 "content_parts": [{"type": "image_url", "url": data_url}],
245 "sampling_params": {"max_tokens": 8, "temperature": 0.0},
246 "stream": True,
247 + "return_token_ids": True,
248 },
249 ) as resp:
250 resp.raise_for_status()
@@ -236,3 +254,9 @@ async def test_content_parts_streaming(client, test_image):
254 chunks.append(json_mod.loads(line[6:]))
255
256 assert len(chunks) > 0
257 + prompt_metadata_chunks = [
258 + chunk for chunk in chunks if chunk.get("prompt_token_ids") is not None
259 + ]
260 + assert len(prompt_metadata_chunks) == 1
261 + assert prompt_metadata_chunks[0]["prompt_token_ids"] == expanded
262 + assert prompt_metadata_chunks[0]["mm_placeholders"] == mm_placeholders
vllm/entrypoints/scale_out/token_in_token_out/mm_features.py
+20 -8
@@ -117,6 +117,24 @@ def mm_kwargs_from_features(
117 return mm_kwargs
118
119
120 +def placeholder_ranges_from_engine_input(
121 + engine_input: EngineInput,
122 +) -> dict[str, list[PlaceholderRangeInfo]] | None:
123 + """Return per-modality placeholder ranges, or ``None`` for text-only prompts."""
124 + if engine_input.get("type") != "multimodal":
125 + return None
126 +
127 + raw_placeholders: MultiModalPlaceholders = cast(MultiModalInput, engine_input)[
128 + "mm_placeholders"
129 + ]
130 + return {
131 + modality: [
132 + PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
133 + ]
134 + for modality, ranges in raw_placeholders.items()
135 + }
136 +
137 +
138 def extract_mm_features(
139 engine_input: EngineInput,
140 *,
@@ -134,14 +152,8 @@ def extract_mm_features(
152
153 mm_engine_input = cast(MultiModalInput, engine_input)
154 mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
137 - raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
138 -
139 - mm_placeholders = {
140 - modality: [
141 - PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
142 - ]
143 - for modality, ranges in raw_placeholders.items()
144 - }
155 + mm_placeholders = placeholder_ranges_from_engine_input(engine_input)
156 + assert mm_placeholders is not None
157
158 kwargs_data: dict[str, list[str | None]] | None = None
159 mm_metadata: dict[str, list[str | None]] | None = None
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+16
@@ -211,6 +211,15 @@ class GenerateRequest(BaseModel):
211
212 model: str | None = None
213
214 + return_token_ids: bool | None = Field(
215 + default=None,
216 + description=(
217 + "If true, return the final prompt token IDs after multimodal "
218 + "placeholder expansion, together with multimodal placeholder ranges. "
219 + "In streaming mode, this metadata is included only in the first chunk."
220 + ),
221 + )
222 +
223 stream: bool | None = False
224 stream_options: StreamOptions | None = None
225 cache_salt: str | None = Field(
@@ -255,6 +264,9 @@ class GenerateRequest(BaseModel):
264 # ``SamplingParams`` instance (e.g. from internal callers that have
265 # already resolved values), in which case all fields are considered set.
266 _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None)
267 + _response_mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = (
268 + PrivateAttr(default=None)
269 + )
270
271 @model_validator(mode="wrap")
272 @classmethod
@@ -354,6 +366,8 @@ class GenerateStreamResponse(BaseModel):
366 )
367 choices: list[GenerateResponseStreamChoice]
368 usage: UsageInfo | None = Field(default=None)
369 + prompt_token_ids: list[int] | None = None
370 + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
371
372
373 class GenerateResponse(BaseModel):
@@ -370,6 +384,8 @@ class GenerateResponse(BaseModel):
384 choices: list[GenerateResponseChoice]
385 usage: UsageInfo | None = Field(default=None)
386 prompt_logprobs: list[dict[int, Logprob] | None] | None = None
387 + prompt_token_ids: list[int] | None = None
388 + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
389
390 kv_transfer_params: dict[str, Any] | None = Field(
391 default=None,
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+35 -3
@@ -44,7 +44,10 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
44 from vllm.utils.collection_utils import as_list
45 from vllm.utils.serial_utils import numpy2base64
46
47 -from .mm_features import mm_kwargs_from_features
47 +from .mm_features import (
48 + mm_kwargs_from_features,
49 + placeholder_ranges_from_engine_input,
50 +)
51 from .protocol import (
52 GenerateRequest,
53 GenerateResponse,
@@ -195,6 +198,14 @@ class ServingTokens(GenerateBaseServing):
198 skip_mm_cache=True,
199 )
200
201 + # Offsets are relative to the decoder prompt, so they are not
202 + # meaningful for encoder-decoder models.
203 + request._response_mm_placeholders = (
204 + placeholder_ranges_from_engine_input(engine_input)
205 + if request.return_token_ids and not self.model_config.is_encoder_decoder
206 + else None
207 + )
208 +
209 # Schedule the request and get the result generator.
210 result_generator: AsyncGenerator[RequestOutput, None] | None = None
211
@@ -359,6 +370,10 @@ class ServingTokens(GenerateBaseServing):
370 choices=choices,
371 usage=usage,
372 prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
373 + prompt_token_ids=(
374 + final_res.prompt_token_ids if request.return_token_ids else None
375 + ),
376 + mm_placeholders=request._response_mm_placeholders,
377 kv_transfer_params=final_res.kv_transfer_params,
378 ec_transfer_params=final_res.ec_transfer_params,
379 )
@@ -395,6 +410,7 @@ class ServingTokens(GenerateBaseServing):
410 num_prompt_tokens = 0
411 num_generated_tokens: list[int] = []
412 first_iteration = True
413 + prompt_token_ids: list[int] | None = None
414 num_cached_tokens = None
415 sampling_params: SamplingParams = request.sampling_params
416
@@ -407,6 +423,8 @@ class ServingTokens(GenerateBaseServing):
423 if first_iteration:
424 if res.prompt_token_ids is not None:
425 num_prompt_tokens = len(res.prompt_token_ids)
426 + if request.return_token_ids:
427 + prompt_token_ids = res.prompt_token_ids
428 if res.encoder_prompt_token_ids is not None:
429 num_prompt_tokens += len(res.encoder_prompt_token_ids)
430 num_cached_tokens = res.num_cached_tokens
@@ -421,7 +439,11 @@ class ServingTokens(GenerateBaseServing):
439 finish_reason = output.finish_reason
440 self._raise_if_error(finish_reason, request_id)
441
424 - if not delta_token_ids:
442 + # Still emit a terminal empty chunk while prompt metadata
443 + # is pending, so zero-token completions deliver it.
444 + if not delta_token_ids and (
445 + finish_reason is None or prompt_token_ids is None
446 + ):
447 continue
448
449 if sampling_params.logprobs is not None:
@@ -453,6 +475,10 @@ class ServingTokens(GenerateBaseServing):
475 )
476 ],
477 )
478 + if prompt_token_ids is not None:
479 + chunk.prompt_token_ids = prompt_token_ids
480 + chunk.mm_placeholders = request._response_mm_placeholders
481 + prompt_token_ids = None
482 if include_continuous_usage:
483 chunk.usage = UsageInfo(
484 prompt_tokens=num_prompt_tokens,
@@ -460,7 +486,13 @@ class ServingTokens(GenerateBaseServing):
486 total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
487 )
488
463 - yield f"data: {chunk.model_dump_json()}\n\n"
489 + # Omit absent prompt metadata (Rust skips None fields).
490 + exclude = {
491 + name
492 + for name in ("prompt_token_ids", "mm_placeholders")
493 + if getattr(chunk, name) is None
494 + }
495 + yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n"
496
497 total_completion_tokens = sum(num_generated_tokens)
498 final_usage_info = UsageInfo(