@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 72d5951d02
Fork 0

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 )