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 3, 03:00 UTC
[Frontend] Add truncation side to OpenAI endpoints (#43260)
Signed-off-by: Rui Zhang <rza21.bc@gmail.com> Signed-off-by: Rui Zhang <rui.zhang@globalrelay.net> Co-authored-by: Rui Zhang <rui.zhang@globalrelay.net>
ruizhang committed
May 21, 2026 at 20:32 UTC
60af5c16ee64ea3c1c573d67d0773a713c87a22e
5 files changed
+147
-6
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+76
@@ -600,6 +600,7 @@ class MockEngine:
600
model_config: MockModelConfig = field(default_factory=MockModelConfig)
601
input_processor: MagicMock = field(default_factory=MagicMock)
602
renderer: MagicMock = field(default_factory=MagicMock)
603
+ errored: bool = False
604
605
606
async def _async_serving_chat_init():
@@ -858,6 +859,50 @@ async def test_serving_chat_truncate_prompt_tokens_max_token_accounting():
859
assert mock_engine.generate.call_args.args[1].max_tokens == 93
860
861
862
+@pytest.mark.asyncio
863
+async def test_serving_chat_truncation_side_controls_prompt_truncation():
864
+ model_config = MockModelConfig()
865
+ model_config.model = MODEL_NAME_SHORT
866
+ model_config.tokenizer = MODEL_NAME_SHORT
867
+ mock_engine = MockEngine(
868
+ model_config=model_config,
869
+ renderer=_build_renderer(model_config),
870
+ )
871
+
872
+ serving_chat = _build_serving_chat(mock_engine)
873
+ messages = [
874
+ {
875
+ "role": "user",
876
+ "content": "Summarize how prompt truncation works in one sentence.",
877
+ }
878
+ ]
879
+
880
+ full_token_ids = await _render_chat_prompt_token_ids(
881
+ serving_chat,
882
+ messages,
883
+ model_name=MODEL_NAME_SHORT,
884
+ )
885
+ assert len(full_token_ids) > 4
886
+
887
+ right_token_ids = await _render_chat_prompt_token_ids(
888
+ serving_chat,
889
+ messages,
890
+ model_name=MODEL_NAME_SHORT,
891
+ truncate_prompt_tokens=4,
892
+ truncation_side="right",
893
+ )
894
+ assert right_token_ids == full_token_ids[:4]
895
+
896
+ left_token_ids = await _render_chat_prompt_token_ids(
897
+ serving_chat,
898
+ messages,
899
+ model_name=MODEL_NAME_SHORT,
900
+ truncate_prompt_tokens=4,
901
+ truncation_side="left",
902
+ )
903
+ assert left_token_ids == full_token_ids[-4:]
904
+
905
+
906
@pytest.mark.asyncio
907
async def test_serving_chat_mistral_token_ids_prompt_is_validated():
908
"""Regression test: when the Mistral tokenizer path returns token IDs
@@ -1111,6 +1156,37 @@ async def test_serving_chat_data_parallel_rank_extraction():
1156
assert mock_engine.generate.call_args.kwargs["data_parallel_rank"] is None
1157
1158
1159
+async def _render_chat_prompt_token_ids(
1160
+ serving_chat: OpenAIServingChat,
1161
+ messages: list[dict[str, str]],
1162
+ *,
1163
+ model_name: str = MODEL_NAME,
1164
+ truncate_prompt_tokens: int | None = None,
1165
+ truncation_side: str | None = None,
1166
+) -> list[int]:
1167
+ request = ChatCompletionRequest(
1168
+ model=model_name,
1169
+ messages=messages,
1170
+ max_tokens=1,
1171
+ temperature=0.0,
1172
+ return_token_ids=True,
1173
+ truncate_prompt_tokens=truncate_prompt_tokens,
1174
+ truncation_side=truncation_side,
1175
+ )
1176
+
1177
+ result = await serving_chat.render_chat_request(request)
1178
+ assert not isinstance(result, ErrorResponse)
1179
+
1180
+ _, engine_inputs = result
1181
+ assert len(engine_inputs) == 1
1182
+
1183
+ prompt_token_ids = serving_chat._extract_prompt_components(
1184
+ engine_inputs[0]
1185
+ ).token_ids
1186
+ assert prompt_token_ids is not None
1187
+ return prompt_token_ids
1188
+
1189
+
1190
class TestServingChatWithHarmony:
1191
"""
1192
These tests ensure Chat Completion requests are being properly converted into
tests/entrypoints/openai/completion/test_completion.py
+37
@@ -73,6 +73,43 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) ->
73
assert completion.choices[0].prompt_logprobs is None
74
75
76
+@pytest.mark.asyncio
77
+@pytest.mark.parametrize(
78
+ "model_name",
79
+ [MODEL_NAME],
80
+)
81
+async def test_completion_truncation_side_controls_prompt_truncation(
82
+ client: openai.AsyncOpenAI, model_name: str
83
+) -> None:
84
+ prompt_token_ids = list(range(8))
85
+
86
+ right_completion = await client.completions.create(
87
+ model=model_name,
88
+ prompt=prompt_token_ids,
89
+ max_tokens=1,
90
+ temperature=0.0,
91
+ extra_body={
92
+ "return_token_ids": True,
93
+ "truncate_prompt_tokens": 4,
94
+ "truncation_side": "right",
95
+ },
96
+ )
97
+ assert right_completion.choices[0].prompt_token_ids == prompt_token_ids[:4]
98
+
99
+ left_completion = await client.completions.create(
100
+ model=model_name,
101
+ prompt=prompt_token_ids,
102
+ max_tokens=1,
103
+ temperature=0.0,
104
+ extra_body={
105
+ "return_token_ids": True,
106
+ "truncate_prompt_tokens": 4,
107
+ "truncation_side": "left",
108
+ },
109
+ )
110
+ assert left_completion.choices[0].prompt_token_ids == prompt_token_ids[-4:]
111
+
112
+
113
@pytest.mark.asyncio
114
@pytest.mark.parametrize(
115
"model_name",
vllm/entrypoints/openai/chat_completion/protocol.py
+9
@@ -245,6 +245,14 @@ class ChatCompletionRequest(OpenAIBaseModel):
245
skip_special_tokens: bool = True
246
spaces_between_special_tokens: bool = True
247
truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_INT64_MAX)] | None = None
248
+ truncation_side: Literal["left", "right"] | None = Field(
249
+ default=None,
250
+ description=(
251
+ "Which side to truncate from when truncate_prompt_tokens is active. "
252
+ "'right' keeps the first N tokens. "
253
+ "'left' keeps the last N tokens."
254
+ ),
255
+ )
256
prompt_logprobs: int | None = None
257
allowed_token_ids: list[int] | None = None
258
bad_words: list[str] = Field(default_factory=list)
@@ -491,6 +499,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
499
max_total_tokens=model_config.max_model_len,
500
max_output_tokens=max_output_tokens or 0,
501
truncate_prompt_tokens=self.truncate_prompt_tokens,
502
+ truncation_side=self.truncation_side,
503
add_special_tokens=self.add_special_tokens,
504
needs_detokenization=bool(self.echo and not self.return_token_ids),
505
max_total_tokens_param="max_model_len",
vllm/entrypoints/openai/completion/protocol.py
+9
@@ -79,6 +79,14 @@ class CompletionRequest(OpenAIBaseModel):
79
skip_special_tokens: bool = True
80
spaces_between_special_tokens: bool = True
81
truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_INT64_MAX)] | None = None
82
+ truncation_side: Literal["left", "right"] | None = Field(
83
+ default=None,
84
+ description=(
85
+ "Which side to truncate from when truncate_prompt_tokens is active. "
86
+ "'right' keeps the first N tokens. "
87
+ "'left' keeps the last N tokens."
88
+ ),
89
+ )
90
allowed_token_ids: list[int] | None = None
91
prompt_logprobs: int | None = None
92
# --8<-- [end:completion-sampling-params]
@@ -192,6 +200,7 @@ class CompletionRequest(OpenAIBaseModel):
200
max_total_tokens=model_config.max_model_len,
201
max_output_tokens=self.max_tokens or 0,
202
truncate_prompt_tokens=self.truncate_prompt_tokens,
203
+ truncation_side=self.truncation_side,
204
add_special_tokens=self.add_special_tokens,
205
needs_detokenization=bool(self.echo and not self.return_token_ids),
206
max_total_tokens_param="max_model_len",
vllm/renderers/params.py
+16
-6
@@ -198,6 +198,13 @@ class TokenizeParams:
198
max_input_tokens = self.max_input_tokens
199
truncate_prompt_tokens = self.truncate_prompt_tokens
200
201
+ if self.truncation_side not in (None, "left", "right"):
202
+ raise VLLMValidationError(
203
+ "`truncation_side` must be either 'left' or 'right'.",
204
+ parameter="truncation_side",
205
+ value=self.truncation_side,
206
+ )
207
+
208
if (
209
max_output_tokens is not None
210
and max_total_tokens is not None
@@ -234,6 +241,9 @@ class TokenizeParams:
241
truncate_prompt_tokens = tokenization_kwargs.pop(
242
"truncate_prompt_tokens", self.truncate_prompt_tokens
243
)
244
+ truncation_side = tokenization_kwargs.pop(
245
+ "truncation_side", self.truncation_side
246
+ )
247
do_lower_case = tokenization_kwargs.pop("do_lower_case", self.do_lower_case)
248
add_special_tokens = tokenization_kwargs.pop(
249
"add_special_tokens", self.add_special_tokens
@@ -279,7 +289,7 @@ class TokenizeParams:
289
),
290
pad_prompt_tokens=pad_prompt_tokens,
291
truncate_prompt_tokens=truncate_prompt_tokens,
282
- truncation_side=self.truncation_side,
292
+ truncation_side=truncation_side,
293
do_lower_case=do_lower_case,
294
add_special_tokens=add_special_tokens,
295
needs_detokenization=needs_detokenization,
@@ -295,11 +305,11 @@ class TokenizeParams:
305
# while still failing `self._token_len_check` as expected by users
306
max_length = self.max_input_tokens + 1
307
298
- # Left-side truncation requires the full token sequence so we can
299
- # slice from the end in _token_truncation. Disable HF-level
300
- # truncation (which would incorrectly truncate from the right for
301
- # pooling models) and let _token_truncation handle it.
302
- if self.truncation_side == "left":
308
+ # Explicit truncation-side overrides require the full token sequence so
309
+ # we can slice from the requested side in _token_truncation. Disable
310
+ # tokenizer-level truncation because generation tokenizers default to
311
+ # left truncation while callers may request right truncation.
312
+ if self.truncation_side is not None and self.truncate_prompt_tokens is not None:
313
return dict(
314
truncation=False,
315
add_special_tokens=self.add_special_tokens,