iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
aoshen02
2026-09-16 15:29:07 -07:00
committed by GitHub
parent 6c3d9816af
commit 0983aef8da
9 changed files with 763 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/server/src/routes/inference/generate.rs
+268 -4
View File
@@ -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<MultiModalPlaceholders>,
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> 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<MultiModalPlaceholders>,
) -> Result<GenerateResponse, ApiError> {
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<MultiModalPlaceholders> {
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<ChatLogProbs, ApiError> {
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<u32>,
finish_reason: Option<FinishReason>,
) -> 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<GenerateOutput>,
return_token_ids: bool,
mm_placeholders: Option<MultiModalPlaceholders>,
) -> Vec<GenerateStreamResponse> {
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<u32>| 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");
}
rust/src/server/src/routes/inference/generate/convert.rs
+4
View File
@@ -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,
},
})
}
rust/src/server/src/routes/inference/generate/types.rs
+13
View File
@@ -46,6 +46,7 @@ pub struct GenerateRequest {
pub ec_transfer_params: Option<HashMap<String, Value>>,
/// Raw multimodal input; server resolves media. Mutually exclusive with `features`.
pub content_parts: Option<Vec<MediaContentPart>>,
pub return_token_ids: Option<bool>,
#[serde(flatten)]
pub other: Map<String, Value>,
}
@@ -81,6 +82,8 @@ pub(super) struct GenerateStreamResponse {
pub request_id: String,
pub choices: Vec<GenerateResponseStreamChoice>,
pub usage: Option<Usage>,
pub prompt_token_ids: Option<Vec<u32>>,
pub mm_placeholders: Option<MultiModalPlaceholders>,
}
/// Mirrors the Python vLLM `GenerateResponse` class.
@@ -89,10 +92,20 @@ pub(super) struct GenerateResponse {
pub request_id: String,
pub choices: Vec<GenerateResponseChoice>,
pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
pub prompt_token_ids: Option<Vec<u32>>,
pub mm_placeholders: Option<MultiModalPlaceholders>,
pub kv_transfer_params: Option<Value>,
pub ec_transfer_params: Option<Value>,
}
pub(super) type MultiModalPlaceholders = HashMap<String, Vec<PlaceholderRangeInfo>>;
#[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 {
rust/src/server/src/routes/render.rs
+1
View File
@@ -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)?;
tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+357
View File
@@ -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]."""
tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py
+49 -25
View File
@@ -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
vllm/entrypoints/scale_out/token_in_token_out/mm_features.py
+20 -8
View File
@@ -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
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+16
View File
@@ -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,
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+35 -3
View File
@@ -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(