[#15327][feat] Add per-request priority support to OpenAI chat/completions

tensorrt_llm/serve/openai_protocol.py

@@ -555,6 +555,15 @@ class CompletionRequest(OpenAIBaseModel):
         default=None,
         description=("Parameters for disaggregated serving"),
     )
+    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."),
+    )
     conversation_params: Optional[ConversationParams] = Field(
         default=None,
         description=("Parameters for multi-turn conversation routing"),
@@ -954,6 +963,15 @@ class ChatCompletionRequest(OpenAIBaseModel):
         ("If specified, KV cache will be salted with the provided string "
          "to limit the kv cache reuse on with the requests having the same string."
          ))
+    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 ")

tensorrt_llm/serve/openai_server.py

@@ -35,6 +35,7 @@
 # yapf: disable
 from tensorrt_llm.executor import CppExecutorError
 from tensorrt_llm.executor.postproc_worker import PostprocParams
+from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY
 from tensorrt_llm.inputs import prompt_inputs
 from tensorrt_llm.inputs.data import TokensPrompt
 from tensorrt_llm.inputs.media_io import BaseMediaIO
@@ -1623,6 +1624,8 @@ 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 DEFAULT_REQUEST_PRIORITY,
             )
             asyncio.create_task(self.await_disconnected(raw_request, promise))
             if not self.postproc_worker_enabled:
@@ -1951,7 +1954,9 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
                     lora_request=request.lora_request,
                     disaggregated_params=disaggregated_params,
                     conversation_params=conversation_params,
-                    trace_headers=trace_headers)
+                    trace_headers=trace_headers,
+                    priority=request.priority if request.priority is not None
+                    else DEFAULT_REQUEST_PRIORITY)
                 asyncio.create_task(
                     self.await_disconnected(raw_request, promise))
                 if not self.postproc_worker_enabled:
@@ -2112,6 +2117,8 @@ async def create_streaming_generator(promise: RequestOutput,
                 disaggregated_params=disaggregated_params,
                 conversation_params=conversation_params,
                 trace_headers=trace_headers,
+                priority=request.priority
+                if request.priority is not None else DEFAULT_REQUEST_PRIORITY,
             )
             if not self.postproc_worker_enabled:
                 postproc_args.num_prompt_tokens = len(promise.prompt_token_ids)

tests/unittest/api_stability/references/trtllm_serve_api.yaml

@@ -301,6 +301,12 @@ models:
         default: null
         status: prototype
         required: false
+      priority:
+        kind: extension
+        type: Optional[float]
+        default: null
+        status: prototype
+        required: false
 
   CompletionResponseChoice:
     fields:
@@ -802,6 +808,12 @@ models:
         default: null
         status: beta
         required: false
+      priority:
+        kind: extension
+        type: Optional[float]
+        default: null
+        status: prototype
+        required: false
 
   ChatCompletionResponseChoice:
     fields: