A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
[Frontend] Allow return_tokens_as_token_ids to be passed as a request param (#14066)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
Benjamin Chislett committed
Mar 5, 2025 at 01:30 UTC
32985bed7c88f654b11f919ead34d77e846c32e3
4 files changed
+64
-25
tests/entrypoints/openai/test_return_tokens_as_ids.py
+30
-14
@@ -17,18 +17,28 @@ from .test_completion import MODEL_NAME
17
18
19
@pytest.fixture(scope="module")
20
-def server_with_return_tokens_as_token_ids_flag(
21
- default_server_args): # noqa: F811
22
- args_with_flag = default_server_args + ["--return-tokens-as-token-ids"]
23
- with RemoteOpenAIServer(MODEL_NAME, args_with_flag) as remote_server:
24
- yield remote_server
20
+def server_fixture(request, default_server_args): # noqa: F811
21
+ use_server_flag = request.param
22
+ if use_server_flag:
23
+ args_with_flag = default_server_args + ["--return-tokens-as-token-ids"]
24
+ with RemoteOpenAIServer(MODEL_NAME, args_with_flag) as remote_server:
25
+ yield (remote_server, True)
26
+ else:
27
+ with RemoteOpenAIServer(MODEL_NAME,
28
+ default_server_args) as remote_server:
29
+ yield (remote_server, False)
30
31
32
@pytest.mark.asyncio
33
+@pytest.mark.parametrize("server_fixture", [True, False], indirect=True)
34
async def test_completion_return_tokens_as_token_ids_completion(
29
- server_with_return_tokens_as_token_ids_flag):
30
- async with server_with_return_tokens_as_token_ids_flag.get_async_client(
31
- ) as client:
35
+ server_fixture):
36
+ server, use_server_flag = server_fixture
37
+ request_args = {}
38
+ if not use_server_flag:
39
+ request_args["return_tokens_as_token_ids"] = True
40
+
41
+ async with server.get_async_client() as client:
42
43
completion = await client.completions.create(
44
model=MODEL_NAME,
@@ -39,7 +49,8 @@ async def test_completion_return_tokens_as_token_ids_completion(
49
echo=True,
50
temperature=0,
51
max_tokens=10,
42
- logprobs=1)
52
+ logprobs=1,
53
+ extra_body=request_args)
54
55
text = completion.choices[0].text
56
token_strs = completion.choices[0].logprobs.tokens
@@ -60,10 +71,14 @@ async def test_completion_return_tokens_as_token_ids_completion(
71
72
73
@pytest.mark.asyncio
63
-async def test_chat_return_tokens_as_token_ids_completion(
64
- server_with_return_tokens_as_token_ids_flag):
65
- async with server_with_return_tokens_as_token_ids_flag.get_async_client(
66
- ) as client:
74
+@pytest.mark.parametrize("server_fixture", [True, False], indirect=True)
75
+async def test_chat_return_tokens_as_token_ids_completion(server_fixture):
76
+ server, use_server_flag = server_fixture
77
+ request_args = {}
78
+ if not use_server_flag:
79
+ request_args["return_tokens_as_token_ids"] = True
80
+
81
+ async with server.get_async_client() as client:
82
response = await client.chat.completions.create(
83
model=MODEL_NAME,
84
# Include Unicode characters to test for dividing a single
@@ -78,7 +93,8 @@ async def test_chat_return_tokens_as_token_ids_completion(
93
}],
94
temperature=0,
95
max_tokens=8,
81
- logprobs=True)
96
+ logprobs=True,
97
+ extra_body=request_args)
98
99
text = response.choices[0].message.content
100
tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
vllm/entrypoints/openai/protocol.py
+12
@@ -369,6 +369,12 @@ class ChatCompletionRequest(OpenAIBaseModel):
369
"arguments. For example: {'qualname': "
370
"'my_module.MyLogitsProcessor', 'args': [1, 2], 'kwargs': "
371
"{'param': 'value'}}."))
372
+ return_tokens_as_token_ids: Optional[bool] = Field(
373
+ default=None,
374
+ description=(
375
+ "If specified with 'logprobs', tokens are represented "
376
+ " as strings of the form 'token_id:{token_id}' so that tokens "
377
+ "that are not JSON-encodable can be identified."))
378
379
# doc: end-chat-completion-extra-params
380
@@ -739,6 +745,12 @@ class CompletionRequest(OpenAIBaseModel):
745
"arguments. For example: {'qualname': "
746
"'my_module.MyLogitsProcessor', 'args': [1, 2], 'kwargs': "
747
"{'param': 'value'}}."))
748
+ return_tokens_as_token_ids: Optional[bool] = Field(
749
+ default=None,
750
+ description=(
751
+ "If specified with 'logprobs', tokens are represented "
752
+ " as strings of the form 'token_id:{token_id}' so that tokens "
753
+ "that are not JSON-encodable can be identified."))
754
755
# doc: end-completion-extra-params
756
vllm/entrypoints/openai/serving_chat.py
+13
-8
@@ -450,6 +450,8 @@ class OpenAIServingChat(OpenAIServing):
450
top_logprobs=output.logprobs,
451
tokenizer=tokenizer,
452
num_output_top_logprobs=request.top_logprobs,
453
+ return_as_token_id=request.
454
+ return_tokens_as_token_ids,
455
)
456
else:
457
logprobs = None
@@ -705,6 +707,7 @@ class OpenAIServingChat(OpenAIServing):
707
top_logprobs=out_logprobs,
708
num_output_top_logprobs=request.top_logprobs,
709
tokenizer=tokenizer,
710
+ return_as_token_id=request.return_tokens_as_token_ids,
711
)
712
else:
713
logprobs = None
@@ -852,13 +855,14 @@ class OpenAIServingChat(OpenAIServing):
855
856
def _get_top_logprobs(
857
self, logprobs: dict[int, Logprob], top_logprobs: Optional[int],
855
- tokenizer: AnyTokenizer) -> list[ChatCompletionLogProb]:
858
+ tokenizer: AnyTokenizer,
859
+ should_return_as_token_id: bool) -> list[ChatCompletionLogProb]:
860
return [
861
ChatCompletionLogProb(token=(token := self._get_decoded_token(
862
p[1],
863
p[0],
864
tokenizer,
861
- return_as_token_id=self.return_tokens_as_token_ids)),
865
+ return_as_token_id=should_return_as_token_id)),
866
logprob=max(p[1].logprob, -9999.0),
867
bytes=list(
868
token.encode("utf-8", errors="replace")))
@@ -872,15 +876,18 @@ class OpenAIServingChat(OpenAIServing):
876
top_logprobs: GenericSequence[Optional[dict[int, Logprob]]],
877
tokenizer: AnyTokenizer,
878
num_output_top_logprobs: Optional[int] = None,
879
+ return_as_token_id: Optional[bool] = None,
880
) -> ChatCompletionLogProbs:
881
"""Create OpenAI-style logprobs."""
882
logprobs_content: list[ChatCompletionLogProbsContent] = []
883
884
+ should_return_as_token_id = return_as_token_id if \
885
+ return_as_token_id is not None else self.return_tokens_as_token_ids
886
for i, token_id in enumerate(token_ids):
887
step_top_logprobs = top_logprobs[i]
888
if step_top_logprobs is None:
889
token = tokenizer.decode(token_id)
883
- if self.return_tokens_as_token_ids:
890
+ if should_return_as_token_id:
891
token = f"token_id:{token_id}"
892
893
logprobs_content.append(
@@ -898,16 +905,14 @@ class OpenAIServingChat(OpenAIServing):
905
step_token,
906
token_id,
907
tokenizer,
901
- self.return_tokens_as_token_ids,
908
+ should_return_as_token_id,
909
),
910
logprob=max(step_token.logprob, -9999.0),
911
bytes=None if step_decoded is None else list(
912
step_decoded.encode("utf-8", errors="replace")),
913
top_logprobs=self._get_top_logprobs(
907
- step_top_logprobs,
908
- num_output_top_logprobs,
909
- tokenizer,
910
- ),
914
+ step_top_logprobs, num_output_top_logprobs,
915
+ tokenizer, should_return_as_token_id),
916
))
917
918
return ChatCompletionLogProbs(content=logprobs_content)
vllm/entrypoints/openai/serving_completion.py
+9
-3
@@ -316,6 +316,8 @@ class OpenAIServingCompletion(OpenAIServing):
316
num_output_top_logprobs=request.logprobs,
317
tokenizer=tokenizer,
318
initial_text_offset=previous_text_lens[i],
319
+ return_as_token_id=request.
320
+ return_tokens_as_token_ids,
321
)
322
else:
323
logprobs = None
@@ -436,6 +438,7 @@ class OpenAIServingCompletion(OpenAIServing):
438
top_logprobs=out_logprobs,
439
tokenizer=tokenizer,
440
num_output_top_logprobs=request.logprobs,
441
+ return_as_token_id=request.return_tokens_as_token_ids,
442
)
443
else:
444
logprobs = None
@@ -477,6 +480,7 @@ class OpenAIServingCompletion(OpenAIServing):
480
num_output_top_logprobs: int,
481
tokenizer: AnyTokenizer,
482
initial_text_offset: int = 0,
483
+ return_as_token_id: Optional[bool] = None,
484
) -> CompletionLogProbs:
485
"""Create logprobs for OpenAI Completion API."""
486
out_text_offset: list[int] = []
@@ -486,11 +490,13 @@ class OpenAIServingCompletion(OpenAIServing):
490
491
last_token_len = 0
492
493
+ should_return_as_token_id = return_as_token_id if \
494
+ return_as_token_id is not None else self.return_tokens_as_token_ids
495
for i, token_id in enumerate(token_ids):
496
step_top_logprobs = top_logprobs[i]
497
if step_top_logprobs is None:
498
token = tokenizer.decode(token_id)
493
- if self.return_tokens_as_token_ids:
499
+ if should_return_as_token_id:
500
token = f"token_id:{token_id}"
501
502
out_tokens.append(token)
@@ -503,7 +509,7 @@ class OpenAIServingCompletion(OpenAIServing):
509
step_token,
510
token_id,
511
tokenizer,
506
- return_as_token_id=self.return_tokens_as_token_ids,
512
+ return_as_token_id=should_return_as_token_id,
513
)
514
token_logprob = max(step_token.logprob, -9999.0)
515
@@ -520,7 +526,7 @@ class OpenAIServingCompletion(OpenAIServing):
526
self._get_decoded_token(top_lp[1],
527
top_lp[0],
528
tokenizer,
523
- return_as_token_id=self.return_tokens_as_token_ids):
529
+ return_as_token_id=should_return_as_token_id):
530
max(top_lp[1].logprob, -9999.0)
531
for i, top_lp in enumerate(step_top_logprobs.items())
532
if num_output_top_logprobs >= i