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 4, 05:00 UTC
[Bugfix] Account for truncate_prompt_tokens when computing max_tokens (#41800)
Signed-off-by: Viktor Pus <viktorpus@tenstorrent.com>
Viktor Pus committed
May 6, 2026 at 18:10 UTC
d5b31c954d511ecdb486e84c71d52d2a5c28ade4
6 files changed
+68
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+51
@@ -807,6 +807,57 @@ async def test_serving_chat_should_set_correct_max_tokens():
807
assert mock_engine.generate.call_args.args[1].max_tokens == 5
808
809
810
+@pytest.mark.asyncio
811
+async def test_serving_chat_truncate_prompt_tokens_max_token_accounting():
812
+ """When truncate_prompt_tokens is set, max_tokens must be calculated using
813
+ the truncated prompt length, not the original prompt length.
814
+
815
+ Regression: without the fix, get_max_tokens received the untruncated prompt
816
+ length, causing the output budget to be underestimated.
817
+ """
818
+ mock_engine = MagicMock(spec=AsyncLLM)
819
+ mock_engine.errored = False
820
+ mock_engine.model_config = MockModelConfig()
821
+ mock_engine.input_processor = MagicMock()
822
+ mock_engine.renderer = _build_renderer(mock_engine.model_config)
823
+
824
+ serving_chat = _build_serving_chat(mock_engine)
825
+
826
+ # "what is 1+1?" tokenizes to 7 tokens with the test chat template
827
+ # (max_model_len=100 -> max_tokens = 93 without truncation, confirmed by
828
+ # test_serving_chat_should_set_correct_max_tokens above).
829
+ messages = [{"role": "user", "content": "what is 1+1?"}]
830
+
831
+ # Baseline: no truncation -> max_tokens = 100 - 7 = 93.
832
+ req = ChatCompletionRequest(model=MODEL_NAME, messages=messages)
833
+ with suppress(Exception):
834
+ await serving_chat.create_chat_completion(req)
835
+ assert mock_engine.generate.call_args.args[1].max_tokens == 93
836
+
837
+ # With truncate_prompt_tokens=5 (less than 7): the effective prompt length
838
+ # is 5, so max_tokens should be 100 - 5 = 95, not 93.
839
+ req = ChatCompletionRequest(
840
+ model=MODEL_NAME,
841
+ messages=messages,
842
+ truncate_prompt_tokens=5,
843
+ )
844
+ with suppress(Exception):
845
+ await serving_chat.create_chat_completion(req)
846
+ assert mock_engine.generate.call_args.args[1].max_tokens == 95
847
+
848
+ # With truncate_prompt_tokens=-1 (meaning use full max_model_len as the
849
+ # truncation limit, i.e., no practical truncation vs the window): effective
850
+ # length = min(7, 100) = 7 -> max_tokens = 93 again.
851
+ req = ChatCompletionRequest(
852
+ model=MODEL_NAME,
853
+ messages=messages,
854
+ truncate_prompt_tokens=-1,
855
+ )
856
+ with suppress(Exception):
857
+ await serving_chat.create_chat_completion(req)
858
+ assert mock_engine.generate.call_args.args[1].max_tokens == 93
859
+
860
+
861
@pytest.mark.asyncio
862
async def test_serving_chat_mistral_token_ids_prompt_is_validated():
863
"""Regression test: when the Mistral tokenizer path returns token IDs
vllm/entrypoints/openai/chat_completion/serving.py
+1
@@ -289,6 +289,7 @@ class OpenAIServingChat(OpenAIServing):
289
self._extract_prompt_len(engine_input),
290
self.default_sampling_params,
291
self.override_max_tokens,
292
+ truncate_prompt_tokens=request.truncate_prompt_tokens,
293
)
294
295
sampling_params: SamplingParams | BeamSearchParams
vllm/entrypoints/openai/completion/serving.py
+1
@@ -151,6 +151,7 @@ class OpenAIServingCompletion(OpenAIServing):
151
self._extract_prompt_len(engine_input),
152
self.default_sampling_params,
153
self.override_max_tokens,
154
+ truncate_prompt_tokens=request.truncate_prompt_tokens,
155
)
156
157
sampling_params: SamplingParams | BeamSearchParams
vllm/entrypoints/openai/responses/serving.py
+6
@@ -416,6 +416,9 @@ class OpenAIServingResponses(OpenAIServing):
416
self._extract_prompt_len(engine_input),
417
self.default_sampling_params,
418
self.override_max_tokens,
419
+ truncate_prompt_tokens=(
420
+ -1 if request.truncation != "disabled" else None
421
+ ),
422
)
423
424
sampling_params = request.to_sampling_params(
@@ -700,6 +703,9 @@ class OpenAIServingResponses(OpenAIServing):
703
self._extract_prompt_len(engine_input),
704
self.default_sampling_params, # type: ignore
705
self.override_max_tokens, # type: ignore
706
+ truncate_prompt_tokens=(
707
+ -1 if context.request.truncation != "disabled" else None
708
+ ),
709
)
710
711
# OPTIMIZATION
vllm/entrypoints/serve/render/serving.py
+2
@@ -164,6 +164,7 @@ class OpenAIServingRender:
164
input_length,
165
self.default_sampling_params,
166
self.override_max_tokens,
167
+ truncate_prompt_tokens=request.truncate_prompt_tokens,
168
)
169
params = request.to_sampling_params(max_tokens, self.default_sampling_params)
170
@@ -298,6 +299,7 @@ class OpenAIServingRender:
299
input_length,
300
self.default_sampling_params,
301
self.override_max_tokens,
302
+ truncate_prompt_tokens=request.truncate_prompt_tokens,
303
)
304
params = request.to_sampling_params(
305
max_tokens, self.default_sampling_params
vllm/entrypoints/utils.py
+7
@@ -177,7 +177,14 @@ def get_max_tokens(
177
input_length: int,
178
default_sampling_params: dict,
179
override_max_tokens: int | None = None,
180
+ truncate_prompt_tokens: int | None = None,
181
) -> int:
182
+ if truncate_prompt_tokens is not None:
183
+ limit = truncate_prompt_tokens
184
+ input_length = min(
185
+ input_length,
186
+ max_model_len if limit == -1 else limit,
187
+ )
188
if max_model_len < input_length:
189
raise ValueError(
190
f"Input length ({input_length}) exceeds model's maximum "