Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions rust/src/server/src/routes/inference/generate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,36 @@ mod tests {
assert!(response.mm_placeholders.is_none());
}

#[test]
fn stream_chunk_omits_absent_fields() {
// Contract shared with the Python frontend: absent optional fields
// are omitted from stream chunks, not serialized as null.
let chunk = GenerateStreamResponse {
request_id: "raw-1".to_string(),
choices: vec![GenerateResponseStreamChoice {
index: 0,
logprobs: None,
finish_reason: None,
token_ids: vec![1],
}],
usage: None,
prompt_token_ids: None,
mm_placeholders: None,
};

let json = serde_json::to_value(&chunk).expect("serialize chunk");

for key in ["usage", "prompt_token_ids", "mm_placeholders"] {
assert!(json.get(key).is_none(), "{key} should be omitted");
}
for key in ["logprobs", "finish_reason"] {
assert!(
json["choices"][0].get(key).is_none(),
"{key} should be omitted"
);
}
}

#[test]
fn extract_mm_placeholders_returns_none_without_features() {
assert!(extract_mm_placeholders(None).is_none());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,40 @@ async def mock_generate(*args, **kwargs):
assert "mm_placeholders" not in chunk


@pytest.mark.asyncio
async def test_stream_chunks_omit_absent_fields():
"""Absent optional fields are omitted from stream chunks, like Rust."""
engine = _mock_engine()

async def mock_generate(*args, **kwargs):
yield _make_request_output("req-1", token_ids=[10])
yield _make_request_output(
"req-1", token_ids=[20], 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=2),
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 len(data_chunks) == 2
for chunk in data_chunks:
assert "usage" not in chunk
assert "prompt_token_ids" not in chunk
assert "mm_placeholders" not in chunk
assert "logprobs" not in chunk["choices"][0]
assert "finish_reason" not in data_chunks[0]["choices"][0]
assert data_chunks[1]["choices"][0]["finish_reason"] == "stop"


@pytest.mark.asyncio
async def test_stream_text_only_prompt_omits_mm_placeholders_key():
engine = _mock_engine()
Expand Down
10 changes: 3 additions & 7 deletions vllm/entrypoints/scale_out/token_in_token_out/serving.py
Original file line number Diff line number Diff line change
Expand Up @@ -486,13 +486,9 @@ async def serve_tokens_stream_generator(
total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
)

# 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"
# Omit absent fields, like the Rust frontend and the
# final usage chunk below.
yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"

total_completion_tokens = sum(num_generated_tokens)
final_usage_info = UsageInfo(
Expand Down
Loading