From 0983aef8da9b25ae5be72a5567e571a8f3f82963 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 17 Sep 2026 06:29:07 +0800 Subject: [PATCH] [Frontend] Return prompt metadata from /inference/v1/generate (#53187) --- .../server/src/routes/inference/generate.rs | 272 ++++++++++++- .../src/routes/inference/generate/convert.rs | 4 + .../src/routes/inference/generate/types.rs | 13 + rust/src/server/src/routes/render.rs | 1 + .../test_generate_stream.py | 357 ++++++++++++++++++ .../test_serving_multimodal_tokens.py | 74 ++-- .../token_in_token_out/mm_features.py | 28 +- .../scale_out/token_in_token_out/protocol.py | 16 + .../scale_out/token_in_token_out/serving.py | 38 +- 9 files changed, 763 insertions(+), 40 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 308ec3ca3f..dfb7e6cbb1 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -21,6 +21,7 @@ use thiserror_ext::AsReport as _; use tracing::{error, info, trace}; use tracing_futures::Instrument as _; use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs}; +use vllm_engine_core_client::protocol::multimodal::MmFeatureSpec; use vllm_llm::{ CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage, }; @@ -28,7 +29,7 @@ use vllm_llm::{ use self::convert::{ResponseOptions, prepare_generate_request}; use self::types::{ GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, - GenerateStreamResponse, + GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo, }; pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; @@ -64,6 +65,11 @@ pub async fn generate( } else { None }; + let mm_placeholders = body + .return_token_ids + .unwrap_or(false) + .then(|| extract_mm_placeholders(mm_features.as_deref())) + .flatten(); let prepared = match prepare_generate_request(body, &lora_resolution, request_context, mm_features) { @@ -98,6 +104,7 @@ pub async fn generate( prepared.request_id, api_server_options, prepared.options, + mm_placeholders, ); let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span); @@ -120,6 +127,7 @@ pub async fn generate( prepared.request_id, api_server_options, prepared.options, + mm_placeholders, ) { Ok(response) => response, Err(error) => return error.into_response(), @@ -143,19 +151,26 @@ async fn generate_chunk_stream( include_logprobs, // Ignored: raw generate streaming has no prompt-logprobs wire shape. include_prompt_logprobs: _, + return_token_ids, }: ResponseOptions, + mut mm_placeholders: Option, mut y: TryYielder, ) -> Result<(), ApiError> { pin_mut!(stream); let mut prompt_tokens = None; + let mut prompt_token_ids = None; let mut usage = TokenUsage::default(); while let Some(next) = stream.next().await { match next { Ok(output) => { - if prompt_tokens.is_none() { - prompt_tokens = - output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len()); + if prompt_tokens.is_none() + && let Some(info) = output.prompt_info.as_ref() + { + prompt_tokens = Some(info.prompt_token_ids.len()); + if return_token_ids { + prompt_token_ids = Some(info.prompt_token_ids.to_vec()); + } } usage.prompt_token_count = prompt_tokens.unwrap_or_default(); usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count); @@ -195,6 +210,7 @@ async fn generate_chunk_stream( None }; + let prompt_token_ids = prompt_token_ids.take(); y.yield_ok(GenerateStreamResponse { request_id: request_id.clone(), choices: vec![GenerateResponseStreamChoice { @@ -205,6 +221,8 @@ async fn generate_chunk_stream( }], usage: include_continuous_usage .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), + mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()), + prompt_token_ids, }) .await; } @@ -223,6 +241,8 @@ async fn generate_chunk_stream( request_id, choices: Vec::new(), usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)), + prompt_token_ids: None, + mm_placeholders: None, }) .await; } @@ -244,7 +264,9 @@ fn collect_generate( include_continuous_usage: _, include_logprobs, include_prompt_logprobs, + return_token_ids, }: ResponseOptions, + mm_placeholders: Option, ) -> Result { let logprobs = if include_logprobs { let logprobs = collected.logprobs.as_ref().ok_or_else(|| { @@ -292,11 +314,34 @@ fn collect_generate( token_ids: collected.token_ids, }], prompt_logprobs, + prompt_token_ids: return_token_ids.then_some(collected.prompt_token_ids), + mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(), kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, }) } +fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option { + let features = features?; + if features.is_empty() { + return None; + } + + let mut placeholders = MultiModalPlaceholders::new(); + for feature in features { + placeholders.entry(feature.modality.as_str().to_string()).or_default().push( + PlaceholderRangeInfo { + offset: feature.mm_position.offset, + length: feature.mm_position.length, + }, + ); + } + for ranges in placeholders.values_mut() { + ranges.sort_by_key(|range| range.offset); + } + Some(placeholders) +} + fn raw_logprobs_to_openai_chat(logprobs: &Logprobs) -> Result { let content = logprobs .positions @@ -419,6 +464,7 @@ mod tests { use std::sync::Arc; use futures::{TryStreamExt as _, stream}; + use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; use vllm_llm::GeneratePromptInfo; use super::*; @@ -451,6 +497,13 @@ mod tests { }), ]); + let mm_placeholders = MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 1, + length: 4, + }], + )]); let chunks: Vec<_> = generate_chunk_stream( stream, "raw-stream".to_string(), @@ -461,14 +514,20 @@ mod tests { ResponseOptions { include_usage: true, include_continuous_usage: true, + return_token_ids: true, ..Default::default() }, + Some(mm_placeholders.clone()), ) .try_collect() .await .expect("collect chunks"); assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22])); + assert_eq!(chunks[0].mm_placeholders, Some(mm_placeholders)); + assert!(chunks[1].prompt_token_ids.is_none()); + assert!(chunks[1].mm_placeholders.is_none()); assert_eq!( chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens, 2 @@ -499,6 +558,209 @@ mod tests { ); } + fn stream_output( + prompt_token_ids: Option<&[u32]>, + token_ids: Vec, + finish_reason: Option, + ) -> GenerateOutput { + GenerateOutput { + request_id: String::new(), + prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { + prompt_token_ids: Arc::from(ids), + prompt_logprobs: None, + }), + token_ids, + logprobs: None, + finish_reason, + cached_token_count: 0, + kv_transfer_params: None, + ec_transfer_params: None, + } + } + + async fn collect_chunks( + outputs: Vec, + return_token_ids: bool, + mm_placeholders: Option, + ) -> Vec { + generate_chunk_stream( + stream::iter(outputs.into_iter().map(Ok)), + "raw-stream".to_string(), + ApiServerOptions::default(), + ResponseOptions { + return_token_ids, + ..Default::default() + }, + mm_placeholders, + ) + .try_collect() + .await + .expect("collect chunks") + } + + #[tokio::test] + async fn generate_chunk_stream_omits_prompt_metadata_by_default() { + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), vec![33], None), + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())), + ], + false, + Some(MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 1, + length: 4, + }], + )])), + ) + .await; + + assert_eq!(chunks.len(), 2); + assert!(chunks.iter().all(|c| c.prompt_token_ids.is_none())); + assert!(chunks.iter().all(|c| c.mm_placeholders.is_none())); + } + + #[tokio::test] + async fn generate_chunk_stream_sends_prompt_metadata_on_first_emitted_chunk() { + // An empty first output is skipped and must not consume the metadata. + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), Vec::new(), None), + stream_output(Some(&[11, 22]), vec![33], None), + stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())), + ], + true, + None, + ) + .await; + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22])); + assert!(chunks[0].mm_placeholders.is_none()); + assert!(chunks[1].prompt_token_ids.is_none()); + assert!(chunks[1].mm_placeholders.is_none()); + } + + #[test] + fn collect_generate_omits_prompt_metadata_by_default() { + let output = CollectedGenerateOutput { + request_id: "raw-1".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 1, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + }; + + let response = collect_generate( + output, + "raw-1".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + Some(MultiModalPlaceholders::from([( + "image".to_string(), + vec![PlaceholderRangeInfo { + offset: 0, + length: 1, + }], + )])), + ) + .expect("response"); + + assert!(response.prompt_token_ids.is_none()); + assert!(response.mm_placeholders.is_none()); + } + + #[test] + fn extract_mm_placeholders_returns_none_without_features() { + assert!(extract_mm_placeholders(None).is_none()); + assert!(extract_mm_placeholders(Some(&[])).is_none()); + } + + #[test] + fn extract_mm_placeholders_groups_ranges_in_prompt_order() { + let feature = |modality: MmModality, offset: usize, length: usize| MmFeatureSpec { + data: None, + modality, + identifier: format!("{}-{offset}", modality.as_str()), + mm_position: PlaceholderRange { + offset, + length, + is_embed: None, + }, + mm_hash: None, + }; + let features = vec![ + feature(MmModality::Image, 8, 2), + feature(MmModality::Audio, 4, 3), + feature(MmModality::Image, 1, 5), + ]; + + let placeholders = extract_mm_placeholders(Some(&features)).expect("metadata"); + + assert_eq!( + placeholders["image"], + vec![ + PlaceholderRangeInfo { + offset: 1, + length: 5, + }, + PlaceholderRangeInfo { + offset: 8, + length: 2, + }, + ] + ); + assert_eq!( + placeholders["audio"], + vec![PlaceholderRangeInfo { + offset: 4, + length: 3, + }] + ); + } + + #[test] + fn collect_generate_returns_prompt_token_ids_when_requested() { + let output = CollectedGenerateOutput { + request_id: "raw-1".to_string(), + prompt_logprobs: None, + token_ids: vec![30], + logprobs: None, + finish_reason: FinishReason::stop_eos(), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 1, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + }; + + let response = collect_generate( + output, + "raw-1".to_string(), + ApiServerOptions::default(), + ResponseOptions { + return_token_ids: true, + ..Default::default() + }, + None, + ) + .expect("response"); + + assert_eq!(response.prompt_token_ids, Some(vec![10, 20])); + } + #[test] fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() { let output_without_payload = |prompt_token_ids: Vec| CollectedGenerateOutput { @@ -525,6 +787,7 @@ mod tests { include_prompt_logprobs: true, ..Default::default() }, + None, ) .expect("single-token prompt without payload maps to [None]"); let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present"); @@ -539,6 +802,7 @@ mod tests { include_prompt_logprobs: true, ..Default::default() }, + None, ) .expect_err("multi-token prompt without payload is an engine failure"); } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 08c4ce18be..8093a34ca2 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -30,6 +30,8 @@ pub(super) struct ResponseOptions { pub include_logprobs: bool, /// Whether the caller requested top-level prompt logprobs. pub include_prompt_logprobs: bool, + /// Whether the caller requested final prompt token metadata. + pub return_token_ids: bool, } /// Validate and lower one raw generate request into the internal @@ -56,6 +58,7 @@ pub(super) fn prepare_generate_request( .unwrap_or(false); let include_logprobs = request.sampling_params.inner.logprobs.is_some(); let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some(); + let return_token_ids = request.return_token_ids.unwrap_or(false); let mut sampling_params = request.sampling_params.inner; sampling_params.vllm_xargs = merge_kv_transfer_params( sampling_params.vllm_xargs, @@ -93,6 +96,7 @@ pub(super) fn prepare_generate_request( include_continuous_usage, include_logprobs, include_prompt_logprobs, + return_token_ids, }, }) } diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index b9220aa7b7..406ec58569 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -46,6 +46,7 @@ pub struct GenerateRequest { pub ec_transfer_params: Option>, /// Raw multimodal input; server resolves media. Mutually exclusive with `features`. pub content_parts: Option>, + pub return_token_ids: Option, #[serde(flatten)] pub other: Map, } @@ -81,6 +82,8 @@ pub(super) struct GenerateStreamResponse { pub request_id: String, pub choices: Vec, pub usage: Option, + pub prompt_token_ids: Option>, + pub mm_placeholders: Option, } /// Mirrors the Python vLLM `GenerateResponse` class. @@ -89,10 +92,20 @@ pub(super) struct GenerateResponse { pub request_id: String, pub choices: Vec, pub prompt_logprobs: Option>>>, + pub prompt_token_ids: Option>, + pub mm_placeholders: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, } +pub(super) type MultiModalPlaceholders = HashMap>; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub(super) struct PlaceholderRangeInfo { + pub offset: usize, + pub length: usize, +} + /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads. #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateLogprob { diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index a4ae0298e9..2d38b63f49 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -105,6 +105,7 @@ fn lower_render_request( kv_transfer_params: None, ec_transfer_params: None, content_parts: None, + return_token_ids: None, other: Default::default(), }; validate_generate_request(&request, &state.served_model_names)?; diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index ad39958702..d9ae424cc8 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -12,6 +12,9 @@ from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.generate.base.protocol import StreamOptions from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.scale_out.token_in_token_out.mm_features import ( + placeholder_ranges_from_engine_input, +) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, GenerateResponse, @@ -19,6 +22,7 @@ from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.exceptions import GenerationError from vllm.logprobs import Logprob +from vllm.multimodal.inputs import PlaceholderRange from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers import renderer_from_config from vllm.renderers.online_renderer import OnlineRenderer @@ -176,6 +180,29 @@ def _parse_sse_chunks(chunks: list[str]) -> list[Any]: return parsed +def test_placeholder_ranges_from_engine_input(): + engine_input: Any = { + "type": "multimodal", + "mm_placeholders": { + "image": [ + PlaceholderRange(offset=1, length=5), + PlaceholderRange(offset=8, length=2), + ] + }, + } + + placeholders = placeholder_ranges_from_engine_input(engine_input) + + assert placeholders is not None + assert [item.model_dump() for item in placeholders["image"]] == [ + {"offset": 1, "length": 5}, + {"offset": 8, "length": 2}, + ] + + text_only: Any = {"type": "token", "prompt_token_ids": [1, 2, 3]} + assert placeholder_ranges_from_engine_input(text_only) is None + + @pytest.mark.asyncio async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution(): engine = _mock_engine() @@ -294,6 +321,336 @@ async def test_stream_basic(): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +async def test_non_stream_returns_final_prompt_token_ids_when_requested(): + engine = _mock_engine() + final_prompt_token_ids = [1, 2, 2, 2, 3] + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + prompt_token_ids=final_prompt_token_ids, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == final_prompt_token_ids + + +@pytest.mark.asyncio +async def test_stream_returns_prompt_token_ids_only_once_when_requested(): + engine = _mock_engine() + final_prompt_token_ids = [1, 2, 2, 2, 3] + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], prompt_token_ids=final_prompt_token_ids + ) + yield _make_request_output( + "req-1", + token_ids=[20], + finish_reason="stop", + finished=True, + prompt_token_ids=final_prompt_token_ids, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks[0]["prompt_token_ids"] == final_prompt_token_ids + assert "prompt_token_ids" not in data_chunks[1] + + +MM_ENGINE_INPUT: Any = { + "type": "multimodal", + "prompt_token_ids": [1, 2, 2, 2, 3], + "mm_placeholders": {"image": [PlaceholderRange(offset=1, length=3)]}, +} +MM_PLACEHOLDERS = {"image": [{"offset": 1, "length": 3}]} + + +def _build_mm_serving_tokens(engine: AsyncLLM) -> ServingTokens: + serving = _build_serving_tokens(engine) + + async def _fake_preprocess(*args, **kwargs): + return [MM_ENGINE_INPUT] + + serving.online_renderer.preprocess_completion = AsyncMock( + side_effect=_fake_preprocess + ) + return serving + + +@pytest.mark.asyncio +@pytest.mark.parametrize("return_token_ids", [False, True]) +async def test_non_stream_prompt_metadata_requires_return_token_ids( + return_token_ids: bool, +): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=return_token_ids, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + if return_token_ids: + assert response.prompt_token_ids == [1, 2, 3] + assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS + else: + assert response.prompt_token_ids is None + assert response.mm_placeholders is None + assert request._response_mm_placeholders is None + + +@pytest.mark.asyncio +async def test_non_stream_text_only_prompt_has_no_mm_placeholders(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == [1, 2, 3] + assert response.mm_placeholders is None + + +@pytest.mark.asyncio +async def test_non_stream_encoder_decoder_omits_mm_placeholders(): + engine = _mock_engine() + engine.model_config.is_encoder_decoder = True + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + + assert isinstance(response, GenerateResponse) + assert response.prompt_token_ids == [1, 2, 3] + assert response.mm_placeholders is None + + +@pytest.mark.asyncio +async def test_stream_prompt_metadata_lands_on_first_emitted_chunk(): + """An empty first output is skipped; metadata goes on the first real chunk, + once, and never on the sibling chunk of the same iteration (n=2).""" + engine = _mock_engine() + + def _two_outputs(tokens0: list[int], tokens1: list[int], **kwargs): + res = _make_request_output("req-1", token_ids=tokens0, **kwargs) + res.outputs.append( + CompletionOutput( + index=1, + text="", + token_ids=tokens1, + cumulative_logprob=None, + logprobs=None, + finish_reason=kwargs.get("finish_reason"), + ) + ) + return res + + async def mock_generate(*args, **kwargs): + yield _two_outputs([], []) + yield _two_outputs([10], [11]) + yield _two_outputs([20], [21], finish_reason="stop", finished=True) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2, n=2), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert [c["choices"][0]["token_ids"] for c in data_chunks] == [ + [10], + [11], + [20], + [21], + ] + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS + for chunk in data_chunks[1:]: + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + + +@pytest.mark.asyncio +async def test_stream_omits_prompt_metadata_by_default(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks + for chunk in data_chunks: + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + + +@pytest.mark.asyncio +async def test_stream_text_only_prompt_omits_mm_placeholders_key(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert "mm_placeholders" not in data_chunks[0] + + +@pytest.mark.asyncio +async def test_stream_zero_token_completion_still_delivers_prompt_metadata(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_mm_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + return_token_ids=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert len(data_chunks) == 1 + assert data_chunks[0]["choices"][0]["token_ids"] == [] + assert data_chunks[0]["choices"][0]["finish_reason"] == "stop" + assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3] + assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS + + +@pytest.mark.asyncio +async def test_stream_zero_token_completion_emits_no_chunk_by_default(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=1), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + + assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")] + assert parsed[-1] == "[DONE]" + + @pytest.mark.asyncio async def test_stream_error_mid_generation(): """finish_reason='error' mid-stream yields error chunk then [DONE].""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 892d3c8285..1d017bd3cf 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -154,11 +154,28 @@ async def test_render_to_generate_roundtrip(client, test_image): ) -@pytest.mark.asyncio -async def test_content_parts_generates_tokens(client, test_image): - """content_parts with raw media should produce output tokens.""" - data_url = encode_image_url(test_image, format="PNG") +def _collapse_placeholders( + token_ids: list[int], mm_placeholders: dict[str, list[dict[str, int]]] +) -> list[int]: + """Undo placeholder expansion: each expanded range becomes one token. + ``/generate`` with ``content_parts`` expects unexpanded placeholders and + expands them itself, so the roundtrip must start from a collapsed prompt. + """ + ranges = sorted( + (r["offset"], r["length"]) for rs in mm_placeholders.values() for r in rs + ) + collapsed: list[int] = [] + pos = 0 + for offset, length in ranges: + collapsed.extend(token_ids[pos:offset]) + collapsed.append(token_ids[offset]) + pos = offset + length + collapsed.extend(token_ids[pos:]) + return collapsed + + +async def _render_image_prompt(client, data_url: str) -> dict: render_resp = await client.post( RENDER_ENDPOINT, json={ @@ -175,14 +192,26 @@ async def test_content_parts_generates_tokens(client, test_image): }, ) render_resp.raise_for_status() - token_ids = render_resp.json()["token_ids"] + return render_resp.json() + + +@pytest.mark.asyncio +async def test_content_parts_generates_tokens(client, test_image): + """Unexpanded prompt + content_parts must expand exactly like /render.""" + data_url = encode_image_url(test_image, format="PNG") + render_data = await _render_image_prompt(client, data_url) + expanded = render_data["token_ids"] + mm_placeholders = render_data["features"]["mm_placeholders"] + unexpanded = _collapse_placeholders(expanded, mm_placeholders) + assert len(unexpanded) < len(expanded) gen_resp = await client.post( GEN_ENDPOINT, json={ - "token_ids": token_ids, + "token_ids": unexpanded, "content_parts": [{"type": "image_url", "url": data_url}], "sampling_params": {"max_tokens": 10, "temperature": 0.0}, + "return_token_ids": True, }, ) gen_resp.raise_for_status() @@ -192,6 +221,8 @@ async def test_content_parts_generates_tokens(client, test_image): choice = gen_data["choices"][0] assert "token_ids" in choice assert len(choice["token_ids"]) > 0 + assert gen_data["prompt_token_ids"] == expanded + assert gen_data["mm_placeholders"] == mm_placeholders @pytest.mark.asyncio @@ -200,33 +231,20 @@ async def test_content_parts_streaming(client, test_image): import json as json_mod data_url = encode_image_url(test_image, format="PNG") - - render_resp = await client.post( - RENDER_ENDPOINT, - json={ - "model": MODEL_NAME, - "messages": [ - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": {"url": data_url}}, - {"type": "text", "text": "Describe this."}, - ], - } - ], - }, - ) - render_resp.raise_for_status() - token_ids = render_resp.json()["token_ids"] + render_data = await _render_image_prompt(client, data_url) + expanded = render_data["token_ids"] + mm_placeholders = render_data["features"]["mm_placeholders"] + unexpanded = _collapse_placeholders(expanded, mm_placeholders) async with client.stream( "POST", GEN_ENDPOINT, json={ - "token_ids": token_ids, + "token_ids": unexpanded, "content_parts": [{"type": "image_url", "url": data_url}], "sampling_params": {"max_tokens": 8, "temperature": 0.0}, "stream": True, + "return_token_ids": True, }, ) as resp: resp.raise_for_status() @@ -236,3 +254,9 @@ async def test_content_parts_streaming(client, test_image): chunks.append(json_mod.loads(line[6:])) assert len(chunks) > 0 + prompt_metadata_chunks = [ + chunk for chunk in chunks if chunk.get("prompt_token_ids") is not None + ] + assert len(prompt_metadata_chunks) == 1 + assert prompt_metadata_chunks[0]["prompt_token_ids"] == expanded + assert prompt_metadata_chunks[0]["mm_placeholders"] == mm_placeholders diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py index d7f89e895e..62d006783a 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py @@ -117,6 +117,24 @@ def mm_kwargs_from_features( return mm_kwargs +def placeholder_ranges_from_engine_input( + engine_input: EngineInput, +) -> dict[str, list[PlaceholderRangeInfo]] | None: + """Return per-modality placeholder ranges, or ``None`` for text-only prompts.""" + if engine_input.get("type") != "multimodal": + return None + + raw_placeholders: MultiModalPlaceholders = cast(MultiModalInput, engine_input)[ + "mm_placeholders" + ] + return { + modality: [ + PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges + ] + for modality, ranges in raw_placeholders.items() + } + + def extract_mm_features( engine_input: EngineInput, *, @@ -134,14 +152,8 @@ def extract_mm_features( mm_engine_input = cast(MultiModalInput, engine_input) mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"] - raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"] - - mm_placeholders = { - modality: [ - PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges - ] - for modality, ranges in raw_placeholders.items() - } + mm_placeholders = placeholder_ranges_from_engine_input(engine_input) + assert mm_placeholders is not None kwargs_data: dict[str, list[str | None]] | None = None mm_metadata: dict[str, list[str | None]] | None = None diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index aec8431b47..be8130e314 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -211,6 +211,15 @@ class GenerateRequest(BaseModel): model: str | None = None + return_token_ids: bool | None = Field( + default=None, + description=( + "If true, return the final prompt token IDs after multimodal " + "placeholder expansion, together with multimodal placeholder ranges. " + "In streaming mode, this metadata is included only in the first chunk." + ), + ) + stream: bool | None = False stream_options: StreamOptions | None = None cache_salt: str | None = Field( @@ -255,6 +264,9 @@ class GenerateRequest(BaseModel): # ``SamplingParams`` instance (e.g. from internal callers that have # already resolved values), in which case all fields are considered set. _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None) + _response_mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = ( + PrivateAttr(default=None) + ) @model_validator(mode="wrap") @classmethod @@ -354,6 +366,8 @@ class GenerateStreamResponse(BaseModel): ) choices: list[GenerateResponseStreamChoice] usage: UsageInfo | None = Field(default=None) + prompt_token_ids: list[int] | None = None + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None class GenerateResponse(BaseModel): @@ -370,6 +384,8 @@ class GenerateResponse(BaseModel): choices: list[GenerateResponseChoice] usage: UsageInfo | None = Field(default=None) prompt_logprobs: list[dict[int, Logprob] | None] | None = None + prompt_token_ids: list[int] | None = None + mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None kv_transfer_params: dict[str, Any] | None = Field( default=None, diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index d3725e3cf7..a8a2804a06 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -44,7 +44,10 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils.collection_utils import as_list from vllm.utils.serial_utils import numpy2base64 -from .mm_features import mm_kwargs_from_features +from .mm_features import ( + mm_kwargs_from_features, + placeholder_ranges_from_engine_input, +) from .protocol import ( GenerateRequest, GenerateResponse, @@ -195,6 +198,14 @@ class ServingTokens(GenerateBaseServing): skip_mm_cache=True, ) + # Offsets are relative to the decoder prompt, so they are not + # meaningful for encoder-decoder models. + request._response_mm_placeholders = ( + placeholder_ranges_from_engine_input(engine_input) + if request.return_token_ids and not self.model_config.is_encoder_decoder + else None + ) + # Schedule the request and get the result generator. result_generator: AsyncGenerator[RequestOutput, None] | None = None @@ -359,6 +370,10 @@ class ServingTokens(GenerateBaseServing): choices=choices, usage=usage, prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), + prompt_token_ids=( + final_res.prompt_token_ids if request.return_token_ids else None + ), + mm_placeholders=request._response_mm_placeholders, kv_transfer_params=final_res.kv_transfer_params, ec_transfer_params=final_res.ec_transfer_params, ) @@ -395,6 +410,7 @@ class ServingTokens(GenerateBaseServing): num_prompt_tokens = 0 num_generated_tokens: list[int] = [] first_iteration = True + prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params @@ -407,6 +423,8 @@ class ServingTokens(GenerateBaseServing): if first_iteration: if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) + if request.return_token_ids: + prompt_token_ids = res.prompt_token_ids if res.encoder_prompt_token_ids is not None: num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens @@ -421,7 +439,11 @@ class ServingTokens(GenerateBaseServing): finish_reason = output.finish_reason self._raise_if_error(finish_reason, request_id) - if not delta_token_ids: + # Still emit a terminal empty chunk while prompt metadata + # is pending, so zero-token completions deliver it. + if not delta_token_ids and ( + finish_reason is None or prompt_token_ids is None + ): continue if sampling_params.logprobs is not None: @@ -453,6 +475,10 @@ class ServingTokens(GenerateBaseServing): ) ], ) + if prompt_token_ids is not None: + chunk.prompt_token_ids = prompt_token_ids + chunk.mm_placeholders = request._response_mm_placeholders + prompt_token_ids = None if include_continuous_usage: chunk.usage = UsageInfo( prompt_tokens=num_prompt_tokens, @@ -460,7 +486,13 @@ class ServingTokens(GenerateBaseServing): total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - yield f"data: {chunk.model_dump_json()}\n\n" + # Omit absent prompt metadata (Rust skips None fields). + exclude = { + name + for name in ("prompt_token_ids", "mm_placeholders") + if getattr(chunk, name) is None + } + yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo(