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 8, 05:00 UTC
[Bugfix] Treat generation_config max_tokens as default not ceiling (#34063)
Signed-off-by: almogtavor <almogtavor@gmail.com>
Almog Tavor committed
Feb 16, 2026 at 17:58 UTC
72d5951d02f2e76228b162fcb63068877850b724
7 files changed
+157
-17
tests/entrypoints/openai/test_serving_chat.py
+50
-10
@@ -526,6 +526,7 @@ class MockModelConfig:
526
allowed_media_domains: list[str] | None = None
527
encoder_config = None
528
generation_config: str = "auto"
529
+ override_generation_config: dict[str, Any] = field(default_factory=dict)
530
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
531
skip_tokenizer_init: bool = False
532
is_encoder_decoder: bool = False
@@ -651,12 +652,10 @@ async def test_serving_chat_should_set_correct_max_tokens():
652
653
assert mock_engine.generate.call_args.args[1].max_tokens == 10
654
654
- # Setting server's max_tokens in the generation_config.json
655
- # lower than context_window - prompt_tokens
655
+ # Model author's generation_config.json sets max_tokens (auto, no override)
656
+ # — should act as fallback only, not ceiling
657
mock_model_config = MockModelConfig()
657
- mock_model_config.diff_sampling_param = {
658
- "max_tokens": 10 # Setting server-side max_tokens limit
659
- }
658
+ mock_model_config.diff_sampling_param = {"max_tokens": 10}
659
660
# Reinitialize the engine with new settings
661
mock_engine = MagicMock(spec=AsyncLLM)
@@ -680,13 +679,14 @@ async def test_serving_chat_should_set_correct_max_tokens():
679
680
assert mock_engine.generate.call_args.args[1].max_tokens == 10
681
683
- # Test Case 2: Request's max_tokens set higher than server accepts
682
+ # Test Case 2: Request's max_tokens set higher than generation_config
683
+ # default so request-provided max_tokens takes precedence
684
req.max_tokens = 15
685
686
with suppress(Exception):
687
await serving_chat.create_chat_completion(req)
688
689
- assert mock_engine.generate.call_args.args[1].max_tokens == 10
689
+ assert mock_engine.generate.call_args.args[1].max_tokens == 15
690
691
# Test Case 3: Request's max_tokens set lower than server accepts
692
req.max_tokens = 5
@@ -696,12 +696,52 @@ async def test_serving_chat_should_set_correct_max_tokens():
696
697
assert mock_engine.generate.call_args.args[1].max_tokens == 5
698
699
+ # User explicitly sets max_tokens via --override-generation-config
700
+ # — should act as a ceiling
701
+ mock_model_config = MockModelConfig()
702
+ mock_model_config.diff_sampling_param = {"max_tokens": 10}
703
+ mock_model_config.override_generation_config = {"max_new_tokens": 10}
704
+
705
+ mock_engine = MagicMock(spec=AsyncLLM)
706
+ mock_engine.errored = False
707
+ mock_engine.model_config = mock_model_config
708
+ mock_engine.input_processor = MagicMock()
709
+ mock_engine.io_processor = MagicMock()
710
+ mock_engine.renderer = _build_renderer(mock_engine.model_config)
711
+
712
+ serving_chat = _build_serving_chat(mock_engine)
713
+
714
+ # Test Case 3.1: No max_tokens — uses override as default
715
+ req = ChatCompletionRequest(
716
+ model=MODEL_NAME,
717
+ messages=[{"role": "user", "content": "what is 1+1?"}],
718
+ )
719
+
720
+ with suppress(Exception):
721
+ await serving_chat.create_chat_completion(req)
722
+
723
+ assert mock_engine.generate.call_args.args[1].max_tokens == 10
724
+
725
+ # Test Case 3.2: Request max_tokens higher — capped by user ceiling from override
726
+ req.max_tokens = 15
727
+
728
+ with suppress(Exception):
729
+ await serving_chat.create_chat_completion(req)
730
+
731
+ assert mock_engine.generate.call_args.args[1].max_tokens == 10
732
+
733
+ # Test Case 3.3: Request max_tokens lower — respected
734
+ req.max_tokens = 5
735
+
736
+ with suppress(Exception):
737
+ await serving_chat.create_chat_completion(req)
738
+
739
+ assert mock_engine.generate.call_args.args[1].max_tokens == 5
740
+
741
# Setting server's max_tokens in the generation_config.json
742
# higher than context_window - prompt_tokens
743
mock_model_config = MockModelConfig()
702
- mock_model_config.diff_sampling_param = {
703
- "max_tokens": 200 # Setting server-side max_tokens limit
704
- }
744
+ mock_model_config.diff_sampling_param = {"max_tokens": 200}
745
746
# Reinitialize the engine with new settings
747
mock_engine = MagicMock(spec=AsyncLLM)
tests/entrypoints/test_utils.py
+73
-1
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from vllm.entrypoints.utils import sanitize_message
3
+
4
+from vllm.entrypoints.utils import get_max_tokens, sanitize_message
5
6
7
def test_sanitize_message():
@@ -8,3 +9,74 @@ def test_sanitize_message():
9
sanitize_message("<_io.BytesIO object at 0x7a95e299e750>")
10
== "<_io.BytesIO object>"
11
)
12
+
13
+
14
+class TestGetMaxTokens:
15
+ """Tests for get_max_tokens() to ensure generation_config's max_tokens
16
+ acts as a default when from model author, and as a ceiling when
17
+ explicitly set by the user."""
18
+
19
+ def test_default_sampling_params_used_when_no_request_max_tokens(self):
20
+ """When user doesn't specify max_tokens, generation_config default
21
+ should apply."""
22
+ result = get_max_tokens(
23
+ max_model_len=24000,
24
+ max_tokens=None,
25
+ input_length=100,
26
+ default_sampling_params={"max_tokens": 2048},
27
+ )
28
+ assert result == 2048
29
+
30
+ def test_request_max_tokens_not_capped_by_default_sampling_params(self):
31
+ """When user specifies max_tokens in request, model author's
32
+ generation_config max_tokens must NOT cap it (fixes #34005)."""
33
+ result = get_max_tokens(
34
+ max_model_len=24000,
35
+ max_tokens=5000,
36
+ input_length=100,
37
+ default_sampling_params={"max_tokens": 2048},
38
+ )
39
+ assert result == 5000
40
+
41
+ def test_override_max_tokens_caps_request(self):
42
+ """When user explicitly sets max_tokens, it acts as a ceiling."""
43
+ result = get_max_tokens(
44
+ max_model_len=24000,
45
+ max_tokens=5000,
46
+ input_length=100,
47
+ default_sampling_params={"max_tokens": 2048},
48
+ override_max_tokens=2048,
49
+ )
50
+ assert result == 2048
51
+
52
+ def test_override_max_tokens_used_as_default(self):
53
+ """When no request max_tokens, override still applies as default."""
54
+ result = get_max_tokens(
55
+ max_model_len=24000,
56
+ max_tokens=None,
57
+ input_length=100,
58
+ default_sampling_params={"max_tokens": 2048},
59
+ override_max_tokens=2048,
60
+ )
61
+ assert result == 2048
62
+
63
+ def test_max_model_len_still_caps_output(self):
64
+ """max_model_len - input_length is always the hard ceiling."""
65
+ result = get_max_tokens(
66
+ max_model_len=3000,
67
+ max_tokens=5000,
68
+ input_length=100,
69
+ default_sampling_params={"max_tokens": 2048},
70
+ )
71
+ assert result == 2900 # 3000 - 100
72
+
73
+ def test_request_max_tokens_smaller_than_default(self):
74
+ """When user explicitly requests fewer tokens than gen_config default,
75
+ that should be respected."""
76
+ result = get_max_tokens(
77
+ max_model_len=24000,
78
+ max_tokens=512,
79
+ input_length=100,
80
+ default_sampling_params={"max_tokens": 2048},
81
+ )
82
+ assert result == 512
vllm/entrypoints/openai/chat_completion/serving.py
+7
@@ -145,6 +145,12 @@ class OpenAIServingChat(OpenAIServing):
145
self.enable_prompt_tokens_details = enable_prompt_tokens_details
146
self.enable_force_include_usage = enable_force_include_usage
147
self.default_sampling_params = self.model_config.get_diff_sampling_param()
148
+ mc = self.model_config
149
+ self.override_max_tokens = (
150
+ self.default_sampling_params.get("max_tokens")
151
+ if mc.generation_config not in ("auto", "vllm")
152
+ else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
153
+ )
154
self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss"
155
if self.use_harmony:
156
if "stop_token_ids" not in self.default_sampling_params:
@@ -389,6 +395,7 @@ class OpenAIServingChat(OpenAIServing):
395
else request.max_tokens,
396
self._extract_prompt_len(engine_prompt),
397
self.default_sampling_params,
398
+ self.override_max_tokens,
399
)
400
401
sampling_params: SamplingParams | BeamSearchParams
vllm/entrypoints/openai/completion/serving.py
+7
@@ -70,6 +70,12 @@ class OpenAIServingCompletion(OpenAIServing):
70
self.enable_force_include_usage = enable_force_include_usage
71
72
self.default_sampling_params = self.model_config.get_diff_sampling_param()
73
+ mc = self.model_config
74
+ self.override_max_tokens = (
75
+ self.default_sampling_params.get("max_tokens")
76
+ if mc.generation_config not in ("auto", "vllm")
77
+ else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
78
+ )
79
80
async def render_completion_request(
81
self,
@@ -164,6 +170,7 @@ class OpenAIServingCompletion(OpenAIServing):
170
request.max_tokens,
171
self._extract_prompt_len(engine_prompt),
172
self.default_sampling_params,
173
+ self.override_max_tokens,
174
)
175
176
sampling_params: SamplingParams | BeamSearchParams
vllm/entrypoints/openai/engine/serving.py
+1
@@ -1174,6 +1174,7 @@ class OpenAIServing:
1174
context.request.max_output_tokens,
1175
self._extract_prompt_len(engine_prompt),
1176
self.default_sampling_params, # type: ignore
1177
+ self.override_max_tokens, # type: ignore
1178
)
1179
1180
# OPTIMIZATION
vllm/entrypoints/openai/responses/serving.py
+7
@@ -229,6 +229,12 @@ class OpenAIServingResponses(OpenAIServing):
229
self.enable_force_include_usage = enable_force_include_usage
230
231
self.default_sampling_params = self.model_config.get_diff_sampling_param()
232
+ mc = self.model_config
233
+ self.override_max_tokens = (
234
+ self.default_sampling_params.get("max_tokens")
235
+ if mc.generation_config not in ("auto", "vllm")
236
+ else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
237
+ )
238
239
# If False (default), the "store" option is (silently) ignored and the
240
# response is not stored. If True, the response is stored in memory.
@@ -446,6 +452,7 @@ class OpenAIServingResponses(OpenAIServing):
452
request.max_output_tokens,
453
self._extract_prompt_len(engine_prompt),
454
self.default_sampling_params,
455
+ self.override_max_tokens,
456
)
457
458
sampling_params = request.to_sampling_params(
vllm/entrypoints/utils.py
+12
-6
@@ -177,17 +177,23 @@ def get_max_tokens(
177
max_tokens: int | None,
178
input_length: int,
179
default_sampling_params: dict,
180
+ override_max_tokens: int | None = None,
181
) -> int:
181
- default_max_tokens = max_model_len - input_length
182
- max_output_tokens = current_platform.get_max_output_tokens(input_length)
182
+ model_max_tokens = max_model_len - input_length
183
+ platform_max_tokens = current_platform.get_max_output_tokens(input_length)
184
+ fallback_max_tokens = (
185
+ max_tokens
186
+ if max_tokens is not None
187
+ else default_sampling_params.get("max_tokens")
188
+ )
189
190
return min(
191
val
192
for val in (
187
- default_max_tokens,
188
- max_tokens,
189
- max_output_tokens,
190
- default_sampling_params.get("max_tokens"),
193
+ model_max_tokens,
194
+ fallback_max_tokens,
195
+ override_max_tokens,
196
+ platform_max_tokens,
197
)
198
if val is not None
199
)