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, 19:00 UTC
[Bugfix][Frontend] Default max_tokens server-side on /inference/v1/generate (#42329)
Signed-off-by: hallerite <git@hallerite.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hallerite committed
May 13, 2026 at 11:16 UTC
97c4317bf5c6e26441f023dbe7d33464072725d4
5 files changed
+168
-2
tests/entrypoints/openai/test_openai_schema.py
+1
@@ -145,6 +145,7 @@ def test_openapi_stateless(case: Case):
145
("POST", "/v1/chat/completions/batch"): LONG_TIMEOUT_SECONDS,
146
("POST", "/v1/completions"): LONG_TIMEOUT_SECONDS,
147
("POST", "/v1/messages"): LONG_TIMEOUT_SECONDS,
148
+ ("POST", "/inference/v1/generate"): LONG_TIMEOUT_SECONDS,
149
}.get(key, DEFAULT_TIMEOUT_SECONDS)
150
151
# No need to verify SSL certificate for localhost
tests/entrypoints/serve/disagg/test_protocol.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the disagg request/response protocol.
4
+
5
+These tests intentionally avoid spinning up a server — they exercise the
6
+pydantic validators on ``GenerateRequest`` directly so they run fast and
7
+fail loudly if the validator semantics ever drift.
8
+"""
9
+
10
+import json
11
+
12
+from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
13
+from vllm.sampling_params import SamplingParams
14
+
15
+
16
+def _base_payload() -> dict:
17
+ return {"token_ids": [1, 2, 3], "sampling_params": {}}
18
+
19
+
20
+def test_omitted_max_tokens_is_not_provided():
21
+ """Body without ``max_tokens`` must surface as 'not provided' so the
22
+ server can apply its own default instead of the dataclass 16."""
23
+ req = GenerateRequest.model_validate(_base_payload())
24
+ # SamplingParams' dataclass default leaks through the parsed instance —
25
+ # this is exactly the bug the server-side defaulting works around.
26
+ assert req.sampling_params.max_tokens == 16
27
+ assert not req.is_sampling_param_provided("max_tokens")
28
+
29
+
30
+def test_explicit_max_tokens_is_provided():
31
+ """Even when the client picks the same value as the dataclass default,
32
+ it must register as explicitly set so the server won't override it."""
33
+ payload = _base_payload()
34
+ payload["sampling_params"] = {"max_tokens": 16}
35
+ req = GenerateRequest.model_validate(payload)
36
+ assert req.sampling_params.max_tokens == 16
37
+ assert req.is_sampling_param_provided("max_tokens")
38
+
39
+ payload["sampling_params"] = {"max_tokens": 256}
40
+ req = GenerateRequest.model_validate(payload)
41
+ assert req.sampling_params.max_tokens == 256
42
+ assert req.is_sampling_param_provided("max_tokens")
43
+
44
+
45
+def test_other_fields_tracked_independently():
46
+ payload = _base_payload()
47
+ payload["sampling_params"] = {"temperature": 0.0}
48
+ req = GenerateRequest.model_validate(payload)
49
+ assert not req.is_sampling_param_provided("max_tokens")
50
+ assert req.is_sampling_param_provided("temperature")
51
+
52
+
53
+def test_json_roundtrip_preserves_provided_keys():
54
+ payload = _base_payload()
55
+ payload["sampling_params"] = {"temperature": 0.5}
56
+ req = GenerateRequest.model_validate_json(json.dumps(payload))
57
+ assert not req.is_sampling_param_provided("max_tokens")
58
+ assert req.is_sampling_param_provided("temperature")
59
+
60
+
61
+def test_internal_instance_construction_treats_all_as_provided():
62
+ """When internal callers build ``GenerateRequest`` from a pre-resolved
63
+ ``SamplingParams`` instance, every field is considered explicitly set
64
+ so server-side defaulting can't clobber values resolved upstream."""
65
+ sp = SamplingParams(max_tokens=500, temperature=0.0)
66
+ req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp)
67
+ assert req.is_sampling_param_provided("max_tokens")
68
+ assert req.is_sampling_param_provided("temperature")
69
+ # And keys we never touched should also count as provided in this path.
70
+ assert req.is_sampling_param_provided("top_p")
tests/entrypoints/serve/disagg/test_serving_tokens.py
+30
@@ -114,6 +114,36 @@ async def test_generate_endpoint(client):
114
assert "choices" in data
115
116
117
+@pytest.mark.asyncio
118
+async def test_generate_defaults_max_tokens_when_omitted(client):
119
+ """Regression: omitting ``max_tokens`` must not silently cap at 16.
120
+
121
+ ``SamplingParams.max_tokens`` defaults to 16. Before the server-side
122
+ defaulting was wired up, every request that didn't set ``max_tokens``
123
+ truncated mid-generation. The server should now fill it in from
124
+ ``max_model_len - prompt_len`` (matching ``/v1/chat/completions``).
125
+ """
126
+ payload = {
127
+ "model": MODEL_NAME,
128
+ "token_ids": [1, 2, 3],
129
+ "sampling_params": {
130
+ "temperature": 0.0,
131
+ "ignore_eos": True,
132
+ },
133
+ "stream": False,
134
+ }
135
+ resp = await client.post(GEN_ENDPOINT, json=payload)
136
+ resp.raise_for_status()
137
+ data = resp.json()
138
+ completion_tokens = len(data["choices"][0]["token_ids"])
139
+ # max_model_len=1024 in the test fixture, prompt is 3 tokens, so we
140
+ # should get ~1021 tokens of output (capped at max_model_len boundary).
141
+ assert completion_tokens > 16, (
142
+ f"expected server-side default to exceed the legacy 16-token cap, "
143
+ f"got {completion_tokens}"
144
+ )
145
+
146
+
147
@pytest.mark.asyncio
148
async def test_generate_stream(client):
149
payload = {
vllm/entrypoints/serve/disagg/protocol.py
+40
-1
@@ -2,7 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Any
4
5
-from pydantic import BaseModel, Field, field_validator
5
+from pydantic import (
6
+ BaseModel,
7
+ Field,
8
+ PrivateAttr,
9
+ field_validator,
10
+ model_validator,
11
+)
12
13
from vllm.config import ModelConfig
14
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionLogProbs
@@ -108,6 +114,39 @@ class GenerateRequest(BaseModel):
114
description="KVTransfer parameters used for disaggregated serving.",
115
)
116
117
+ # Tracks which keys the caller explicitly set inside ``sampling_params``
118
+ # when the request was parsed from a JSON body. Lets the server tell
119
+ # "client said max_tokens=16" from "client said nothing → dataclass
120
+ # default 16" so it can apply server-side defaulting only in the latter
121
+ # case. ``None`` means the request was constructed with a pre-built
122
+ # ``SamplingParams`` instance (e.g. from internal callers that have
123
+ # already resolved values), in which case all fields are considered set.
124
+ _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None)
125
+
126
+ @model_validator(mode="wrap")
127
+ @classmethod
128
+ def _capture_sampling_params_provided_keys(cls, data: Any, handler):
129
+ provided: set[str] | None = None
130
+ if isinstance(data, dict):
131
+ sp = data.get("sampling_params")
132
+ if isinstance(sp, dict):
133
+ provided = set(sp.keys())
134
+ instance = handler(data)
135
+ instance._sampling_params_provided_keys = provided
136
+ return instance
137
+
138
+ def is_sampling_param_provided(self, name: str) -> bool:
139
+ """Whether the caller explicitly set ``sampling_params.<name>``.
140
+
141
+ For requests parsed from a JSON body, this reflects the raw input
142
+ dict. For requests constructed with a pre-built ``SamplingParams``
143
+ instance, all fields are considered provided so server-side defaults
144
+ do not clobber values already resolved upstream.
145
+ """
146
+ if self._sampling_params_provided_keys is None:
147
+ return True
148
+ return name in self._sampling_params_provided_keys
149
+
150
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
151
return TokenizeParams(
152
max_total_tokens=None,
vllm/entrypoints/serve/disagg/serving.py
+27
-1
@@ -34,7 +34,7 @@ from vllm.entrypoints.serve.disagg.protocol import (
34
GenerateStreamResponse,
35
)
36
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
37
-from vllm.entrypoints.utils import should_include_usage
37
+from vllm.entrypoints.utils import get_max_tokens, should_include_usage
38
from vllm.inputs import EngineInput, mm_input
39
from vllm.logger import init_logger
40
from vllm.logprobs import Logprob
@@ -81,6 +81,18 @@ class ServingTokens(OpenAIServing):
81
"step for incoming requests."
82
)
83
84
+ # Mirrors ``OpenAIServingChat`` so we can apply server-side
85
+ # ``max_tokens`` defaulting when the client omits it. Without this,
86
+ # ``SamplingParams.max_tokens`` falls back to its dataclass default
87
+ # of 16 and silently truncates every generation.
88
+ self.default_sampling_params = self.model_config.get_diff_sampling_param()
89
+ mc = self.model_config
90
+ self.override_max_tokens = (
91
+ self.default_sampling_params.get("max_tokens")
92
+ if mc.generation_config not in ("auto", "vllm")
93
+ else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
94
+ )
95
+
96
async def serve_tokens(
97
self,
98
request: GenerateRequest,
@@ -150,6 +162,20 @@ class ServingTokens(OpenAIServing):
162
# Schedule the request and get the result generator.
163
result_generator: AsyncGenerator[RequestOutput, None] | None = None
164
sampling_params = request.sampling_params
165
+
166
+ # Apply server-side ``max_tokens`` defaulting when the client did
167
+ # not set it, matching the OpenAI-compat endpoints. ``SamplingParams``
168
+ # defaults ``max_tokens`` to 16, which would otherwise silently cap
169
+ # every generation that omits the field.
170
+ if not request.is_sampling_param_provided("max_tokens"):
171
+ sampling_params.max_tokens = get_max_tokens(
172
+ max_model_len=self.model_config.max_model_len,
173
+ max_tokens=None,
174
+ input_length=self._extract_prompt_len(engine_input),
175
+ default_sampling_params=self.default_sampling_params,
176
+ override_max_tokens=self.override_max_tokens,
177
+ )
178
+
179
if self.force_no_detokenize:
180
sampling_params.detokenize = False
181
if request.stream: