diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 8bcf8c6d34..5a2e329708 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -12,8 +12,10 @@ import pytest_asyncio from tests.utils import RemoteLaunchRenderServer from vllm.entrypoints.generate.base.protocol import FunctionCall from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateTokensResponse, + GenerateTokensStreamResponse, ) from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.tokenizers import get_tokenizer @@ -1579,3 +1581,43 @@ def test_derender_chat_keeps_parser_tool_call_id(): (tool_call,) = choice.message.tool_calls assert tool_call.id == "functions.get_weather:1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("return_token_ids", [None, False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_derender_completion_preserves_requested_token_ids( + tokenizer, return_token_ids, stream +): + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + renderer._executor = None + model_config = MagicMock() + model_config.model = MODEL_NAME + model_config.hf_config.model_type = "llama" + derenderer = OnlineDerenderer( + model_config, + renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + ) + token_ids = tokenizer.encode("Hello world", add_special_tokens=False) + request = CompletionRequest( + model=MODEL_NAME, prompt="Hi", return_token_ids=return_token_ids + ) + payload = _make_generate_response(token_ids, finish_reason="length") + if stream: + response, _ = await derenderer.derender_completion_stream( + MODEL_NAME, + GenerateTokensStreamResponse.model_validate(payload), + completion_request=request, + ) + choices = response.choices + else: + choices, _, _ = await derenderer.derender_completion( + [GenerateTokensResponse.model_validate(payload)], + completion_request=request, + ) + assert choices[0].text == tokenizer.decode(token_ids) + assert choices[0].token_ids == (token_ids if return_token_ids else None) diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index f0169d1dba..753dc6084b 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -834,6 +834,9 @@ class OnlineDerenderer: total_completion_tokens = 0 index = 0 top_limit = _completion_top_logprobs_limit(completion_request) + return_token_ids = ( + completion_request is not None and completion_request.return_token_ids + ) for gen, pt, seed_ids_override in zip( generate_responses, prompt_tokens_list, prompt_token_ids_list @@ -877,6 +880,7 @@ class OnlineDerenderer: finish_reason=choice.finish_reason, stop_reason=choice.stop_reason, logprobs=completion_logprobs, + token_ids=choice.token_ids if return_token_ids else None, ) ) total_completion_tokens += len(choice.token_ids) @@ -942,6 +946,9 @@ class OnlineDerenderer: updated_state = state top_limit = _completion_top_logprobs_limit(completion_request) + return_token_ids = ( + completion_request is not None and completion_request.return_token_ids + ) for choice in generate_chunk.choices: delta_tids = choice.token_ids or [] new_text, updated_state = await self._detokenize_delta_async( @@ -982,6 +989,7 @@ class OnlineDerenderer: logprobs=completion_logprobs, finish_reason=choice.finish_reason, stop_reason=choice.stop_reason, + token_ids=choice.token_ids if return_token_ids else None, ) )