feat(api): cap request output tokens

docs/CN/source/tutorial/api_server_args.rst

@@ -145,6 +145,11 @@ PD 分离模式参数
 
     新批次的最大 token 数量，控制预填充批次大小以防止 OOM
 
+.. option:: --max-request-output-tokens
+
+    单请求输出 token 的默认值和硬上限，默认为 ``65536``。
+    请求中更大的 ``max_new_tokens`` 或 ``max_tokens`` 会被截断到该值。
+
 .. option:: --running_max_req_size
 
     同时进行前向推理的最大请求数量，默认为 ``1000``

docs/EN/source/tutorial/api_server_args.rst

@@ -147,6 +147,11 @@ Memory and Batch Processing Parameters
 
     Maximum token count for new batches, controls prefill batch size to prevent OOM
 
+.. option:: --max-request-output-tokens
+
+    Default and hard limit for output tokens per request, default is ``65536``.
+    Requests with a larger ``max_new_tokens`` or ``max_tokens`` value are capped at this limit.
+
 .. option:: --running_max_req_size
 
     Maximum number of requests for simultaneous forward inference, default is ``1000``

lightllm/server/api_cli.py

@@ -174,6 +174,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
         default=None,
         help="max tokens num for new cat batch, it control prefill batch size to Preventing OOM",
     )
+    parser.add_argument(
+        "--max-request-output-tokens",
+        "--max_request_output_tokens",
+        dest="max_request_output_tokens",
+        type=int,
+        default=65536,
+        help="default and hard limit for the number of output tokens per request",
+    )
     parser.add_argument(
         "--eos_id", nargs="+", type=int, default=None, help="eos stop token id, if None, will load from config.json"
     )

lightllm/server/api_models.py

@@ -123,7 +123,7 @@ class CompletionRequest(BaseModel):
     prompt: Union[str, List[str], List[int], List[List[int]]]
     suffix: Optional[str] = None
     max_tokens: Optional[int] = Field(
-        default=65536, deprecated="max_tokens is deprecated, please use max_completion_tokens instead"
+        default=None, deprecated="max_tokens is deprecated, please use max_completion_tokens instead"
     )
     max_completion_tokens: Optional[int] = None
     temperature: Optional[float] = 1.0
@@ -199,7 +199,7 @@ class ChatCompletionRequest(BaseModel):
     stream_options: Optional[StreamOptions] = None
     stop: Optional[Union[str, List[str]]] = None
     max_tokens: Optional[int] = Field(
-        default=65536, deprecated="max_tokens is deprecated, please use max_completion_tokens instead"
+        default=None, deprecated="max_tokens is deprecated, please use max_completion_tokens instead"
     )
     max_completion_tokens: Optional[int] = None
     presence_penalty: Optional[float] = 0.0

lightllm/server/api_start.py

@@ -120,6 +120,8 @@ def _launch_subprocesses(args: StartArgs):
     assert (
         args.mem_fraction > 0 and args.mem_fraction < 1
     ), f"Invalid mem_fraction {args.mem_fraction}, The expected value is between 0 and 1."
+    if args.max_request_output_tokens < 1:
+        raise ValueError("max_request_output_tokens must be a positive integer.")
 
     if args.graph_max_len_in_batch == 0:
         args.graph_max_len_in_batch = args.max_req_total_len

lightllm/server/core/objs/sampling_params.py

@@ -333,7 +333,8 @@ def init(self, tokenizer, **kwargs):
         self.top_k = kwargs.get("top_k", SamplingParams._top_k)
         self.ignore_eos = kwargs.get("ignore_eos", False)
         self.image_max_patch_num = kwargs.get("image_max_patch_num", -1)
-        self.max_new_tokens = kwargs.get("max_new_tokens", 65535)
+        max_request_output_tokens = get_env_start_args().max_request_output_tokens
+        self.max_new_tokens = min(kwargs.get("max_new_tokens", max_request_output_tokens), max_request_output_tokens)
         self.min_new_tokens = kwargs.get("min_new_tokens", 1)
         self.input_penalty = kwargs.get("input_penalty", DEFAULT_INPUT_PENALTY)
         self.group_request_id = kwargs.get("group_request_id", -1)

lightllm/server/core/objs/start_args_type.py

@@ -40,6 +40,7 @@ class StartArgs:
     max_total_token_num: Optional[int] = field(default=None)
     mem_fraction: float = field(default=0.8)
     batch_max_tokens: Optional[int] = field(default=None)
+    max_request_output_tokens: int = field(default=65536)
     eos_id: Optional[List[int]] = field(default=None)
     tool_call_parser: Optional[str] = field(
         default=None,

test/test_api/test_max_request_output_tokens.py

@@ -0,0 +1,18 @@
+from types import SimpleNamespace
+
+import lightllm.server.core.objs.sampling_params as sampling_params_module
+from lightllm.server.core.objs.sampling_params import SamplingParams
+
+
+def test_max_request_output_tokens_is_default_and_hard_limit(monkeypatch):
+    monkeypatch.setattr(
+        sampling_params_module,
+        "get_env_start_args",
+        lambda: SimpleNamespace(max_request_output_tokens=1024),
+    )
+
+    for requested, expected in ((None, 1024), (256, 256), (2048, 1024)):
+        params = SamplingParams()
+        kwargs = {} if requested is None else {"max_new_tokens": requested}
+        params.init(None, **kwargs)
+        assert params.max_new_tokens == expected

test/test_api/test_seed_validation.py

@@ -1,11 +1,23 @@
+from types import SimpleNamespace
+
 import pytest
 from pydantic import ValidationError
 
+import lightllm.server.core.objs.sampling_params as sampling_params_module
 from lightllm.server.api_models import ChatCompletionRequest, CompletionRequest, MAX_SEED
 from lightllm.server.core.objs.py_sampling_params import SamplingParams as PySamplingParams
 from lightllm.server.core.objs.sampling_params import SamplingParams
 
 
+@pytest.fixture(autouse=True)
+def mock_start_args(monkeypatch):
+    monkeypatch.setattr(
+        sampling_params_module,
+        "get_env_start_args",
+        lambda: SimpleNamespace(max_request_output_tokens=65536),
+    )
+
+
 @pytest.mark.parametrize(
     ("request_type", "request_data"),
     [