diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 1c0918fe370e..de1d6ea29d47 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -435,6 +435,15 @@ class CompletionRequest(OpenAIBaseModel): "The time interval in milliseconds to create responses under the streaming mode. " "If not set, the engine-level default is used."), ) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description=( + "Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) # doc: end-completion-extra-params @@ -828,6 +837,15 @@ class ChatCompletionRequest(OpenAIBaseModel): "The time interval in milliseconds to create responses under the streaming mode. " "If not set, the engine-level default is used."), ) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description=( + "Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) agent_hierarchy: Optional[AgentHierarchy] = Field( default=None, description="Agent hierarchy ") diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index f1a9a012de29..c82464b06d4c 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1597,6 +1597,7 @@ async def chat_stream_generator( cache_salt=request.cache_salt, trace_headers=trace_headers, scheduling_params=scheduling_params, + priority=request.priority if request.priority is not None else 0.5, ) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer @@ -1921,7 +1922,8 @@ async def generator_wrapper(generator: AsyncIterator[Any]): streaming=request.stream, lora_request=request.lora_request, disaggregated_params=disaggregated_params, - trace_headers=trace_headers + trace_headers=trace_headers, + priority=request.priority if request.priority is not None else 0.5, ) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer