feat: add priority-aware request scheduling

README.md

@@ -169,6 +169,24 @@
         > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录
 
   - 试验中功能
+    - 请求优先级调度
+
+      OpenAI 兼容接口接受可选的整数参数 `priority`，范围为 0 到 10，数值越大越优先。默认值为 0；相同有效优先级的请求保持 FIFO。该功能只调整等待请求进入 Prefill 的顺序，不抢占正在运行的请求。
+
+      ```json
+      {
+        "model": "model-name",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "priority": 8
+      }
+      ```
+
+      为避免低优先级请求饥饿，等待请求默认每 5 秒提升一级有效优先级。可在启动服务时调整 Aging 间隔：
+
+      ```bash
+      python python/infinilm/server/inference_server.py --model=<model-path> --priority-aging-interval=10
+      ```
+
     - Warm Up
       ```bash
       python examples/bench.py --device nvidia --model=<model-path> --warmup

python/infinilm/base_config.py

@@ -5,6 +5,7 @@
 import shutil
 import warnings
 
+from infinilm.config.engine_config import DEFAULT_PRIORITY_AGING_INTERVAL
 from infinilm.moe_config import MOE_EP_BACKEND_HELP
 
 
@@ -81,6 +82,7 @@ def __init__(self):
         self.num_blocks = self.args.num_blocks
         self.block_size = self.args.block_size
         self.max_cache_len = self.args.max_cache_len
+        self.priority_aging_interval = self.args.priority_aging_interval
         self.kv_cache_dtype = self.args.kv_cache_dtype
         self.skip_load = self.args.skip_load
         self.weight_load_mode = self.args.weight_load_mode
@@ -283,6 +285,12 @@ def _add_common_args(self):
         self.parser.add_argument(
             "--max-cache-len", type=int, default=4096, help="maximum cache length"
         )
+        self.parser.add_argument(
+            "--priority-aging-interval",
+            type=float,
+            default=DEFAULT_PRIORITY_AGING_INTERVAL,
+            help="seconds before a waiting request gains one priority level",
+        )
         self.parser.add_argument(
             "--kv-cache-dtype",
             type=str,

python/infinilm/config/engine_config.py

@@ -1,8 +1,25 @@
+import math
 from dataclasses import dataclass
 from typing import Optional
 
 from infinilm.config.kv_transfer import KVTransferConfig
 
+DEFAULT_PRIORITY_AGING_INTERVAL = 5.0
+
+
+def validate_priority_aging_interval(interval: float) -> float:
+    """Validate and normalize the priority aging interval."""
+    if (
+        isinstance(interval, bool)
+        or not isinstance(interval, (int, float))
+        or not math.isfinite(interval)
+        or interval <= 0
+    ):
+        raise ValueError(
+            "`priority_aging_interval` must be a finite number greater than zero."
+        )
+    return float(interval)
+
 
 @dataclass
 class EngineConfig:
@@ -27,6 +44,7 @@ class EngineConfig:
         num_blocks: Number of KV cache blocks (only for paged cache).
         block_size: Size of each KV cache block (only for paged cache).
         max_cache_len: Maximum sequence length (only for static cache).
+        priority_aging_interval: Seconds before a waiting request gains one priority level.
         enable_prefix_caching: Whether to reuse KV cache across requests.
         temperature: Default sampling temperature.
         top_p: Default top-p sampling parameter.
@@ -69,6 +87,7 @@ class EngineConfig:
     use_legacy_moe: bool = False
     kv_transfer_config: Optional[KVTransferConfig] = None
     enable_prefix_caching: bool = True
+    priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL
 
     def __post_init__(self) -> None:
         if self.num_draft_tokens < 1:
@@ -84,6 +103,9 @@ def __post_init__(self) -> None:
 
         if self.weight_load_mode not in {"async", "sync"}:
             raise ValueError("weight_load_mode must be either 'async' or 'sync'")
+        self.priority_aging_interval = validate_priority_aging_interval(
+            self.priority_aging_interval
+        )
 
         if (
             self.kv_transfer_config is not None

python/infinilm/llm/llm.py

@@ -16,7 +16,10 @@
 
 import janus
 
-from infinilm.config.engine_config import EngineConfig
+from infinilm.config.engine_config import (
+    DEFAULT_PRIORITY_AGING_INTERVAL,
+    EngineConfig,
+)
 from infinilm.config.kv_transfer import KVTransferConfig
 from infinilm.infer_engine import model_uses_mamba_cache, read_hf_config
 from infinilm.kv_connector import KVConnectorFactory, KVConnectorRole
@@ -70,6 +73,7 @@ def __init__(self, config: EngineConfig):
             self.scheduler = StaticScheduler(
                 max_cache_len=config.max_cache_len,
                 enable_prefix_caching=config.enable_prefix_caching,
+                priority_aging_interval=config.priority_aging_interval,
             )
             logger.info(
                 f"Using Static KV Cache with max_cache_len={config.max_cache_len}"
@@ -110,6 +114,7 @@ def __init__(self, config: EngineConfig):
                 has_mamba_cache=has_mamba_cache,
                 num_mamba_cache_blocks=num_mamba_cache_blocks,
                 enable_prefix_caching=config.enable_prefix_caching,
+                priority_aging_interval=config.priority_aging_interval,
             )
             logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}")
             if has_mamba_cache:
@@ -366,6 +371,7 @@ def __init__(
         skip_load: bool = False,
         use_legacy_moe: bool = False,
         enable_prefix_caching: bool = True,
+        priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL,
     ):
         """Initialize LLM.
 
@@ -387,6 +393,8 @@ def __init__(
             attn_backend: Attention backend to use ('default', 'flash-attn').
             use_mla: Whether to use DeepSeek V2 MLA attention when supported.
             weight_load_mode: Weight loading mode across tensor-parallel workers.
+            priority_aging_interval: Seconds before a waiting request gains one
+                priority level.
         """
         config = EngineConfig(
             model_path=model_path,
@@ -407,6 +415,7 @@ def __init__(
             num_blocks=num_blocks,
             block_size=block_size,
             max_cache_len=max_cache_len,
+            priority_aging_interval=priority_aging_interval,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
@@ -594,6 +603,7 @@ def __init__(
         weight_load_mode: str = "async",
         use_legacy_moe: bool = False,
         enable_prefix_caching: bool = True,
+        priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL,
     ):
         """Initialize AsyncLLMEngine.
 
@@ -619,6 +629,8 @@ def __init__(
             use_mla: Whether to use DeepSeek V2 MLA attention when supported.
             skip_load: Whether to skip loading model weights.
             weight_load_mode: Weight loading mode across tensor-parallel workers.
+            priority_aging_interval: Seconds before a waiting request gains one
+                priority level.
         """
         config = EngineConfig(
             model_path=model_path,
@@ -639,6 +651,7 @@ def __init__(
             num_blocks=num_blocks,
             block_size=block_size,
             max_cache_len=max_cache_len,
+            priority_aging_interval=priority_aging_interval,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
@@ -784,6 +797,7 @@ def add_request(
         request_id: Optional[str] = None,
         # For server use
         request_data: Optional[dict] = None,
+        priority: int = 0,
     ) -> InferenceRequest:
         """Add a request to the engine.
 
@@ -813,6 +827,7 @@ def add_request(
             sampling_params: Sampling parameters.
             request_id: Optional request ID.
             request_data: Optional request data dict (for server use).
+            priority: Request admission priority from 0 to 10.
 
         Returns:
             The created InferenceRequest object.
@@ -878,6 +893,7 @@ def add_request(
             sampling_params=sampling_params,
             eos_token_ids=self.engine.eos_token_ids,
             request_data=request_data,
+            priority=priority,
         )
 
         if request_data and "kv_transfer_params" in request_data:
@@ -897,6 +913,7 @@ def add_chat_request(
         request_id: Optional[str] = None,
         request_data: Optional[dict] = None,
         add_generation_prompt: bool = True,
+        priority: int = 0,
         **kwargs,
     ) -> InferenceRequest:
         """Add a chat request to the engine.
@@ -906,6 +923,7 @@ def add_chat_request(
             sampling_params: Sampling parameters.
             request_id: Optional request ID.
             request_data: Optional request data dict.
+            priority: Request admission priority from 0 to 10.
 
         Returns:
             The created InferenceRequest object.
@@ -918,6 +936,7 @@ def add_chat_request(
             sampling_params=sampling_params,
             request_id=request_id,
             request_data=request_data,
+            priority=priority,
         )
 
     async def stream_request(

python/infinilm/llm/priority_scheduling.py

@@ -0,0 +1,113 @@
+"""Priority ordering helpers for inference request admission."""
+
+import queue
+import time
+from collections import Counter
+from threading import Lock
+from typing import Callable
+
+from infinilm.llm.request import InferenceRequest
+
+
+class PrioritySchedulingStats:
+    """Track cumulative priority admission statistics."""
+
+    def __init__(self) -> None:
+        self.admitted_requests = 0
+        self.aged_admissions = 0
+        self.total_wait_time_seconds = 0.0
+        self.max_wait_time_seconds = 0.0
+        self.admitted_by_priority: Counter[int] = Counter()
+        self._lock = Lock()
+
+    def record_admission(
+        self,
+        request: InferenceRequest,
+        now: float,
+        aging_interval: float,
+    ) -> tuple[float, int, bool]:
+        """Record one admission and return its scheduling details."""
+        wait_time = request_wait_time(request, now)
+        admission_priority = effective_priority(request, now, aging_interval)
+        aged = admission_priority > request.priority
+
+        with self._lock:
+            self.admitted_requests += 1
+            self.aged_admissions += int(aged)
+            self.total_wait_time_seconds += wait_time
+            self.max_wait_time_seconds = max(self.max_wait_time_seconds, wait_time)
+            self.admitted_by_priority[request.priority] += 1
+        return wait_time, admission_priority, aged
+
+    def snapshot(self) -> dict:
+        """Return a JSON-serializable statistics snapshot."""
+        with self._lock:
+            average_wait_time = (
+                self.total_wait_time_seconds / self.admitted_requests
+                if self.admitted_requests
+                else 0.0
+            )
+            return {
+                "admitted_requests": self.admitted_requests,
+                "aged_admissions": self.aged_admissions,
+                "average_wait_time_seconds": average_wait_time,
+                "max_wait_time_seconds": self.max_wait_time_seconds,
+                "admitted_by_priority": dict(sorted(self.admitted_by_priority.items())),
+            }
+
+
+def mark_request_enqueued(
+    request: InferenceRequest,
+    sequence: int,
+    clock: Callable[[], float] = time.monotonic,
+) -> None:
+    """Record stable scheduling metadata on a request's first enqueue."""
+    if request.scheduling_enqueue_time is None:
+        request.scheduling_enqueue_time = clock()
+    if request.scheduling_sequence is None:
+        request.scheduling_sequence = sequence
+
+
+def effective_priority(
+    request: InferenceRequest,
+    now: float,
+    aging_interval: float,
+) -> int:
+    """Return the request priority after applying wait-time aging."""
+    if aging_interval <= 0:
+        raise ValueError("`aging_interval` must be greater than zero.")
+    wait_time = request_wait_time(request, now)
+    return request.priority + int(wait_time // aging_interval)
+
+
+def request_wait_time(request: InferenceRequest, now: float) -> float:
+    """Return how long a request has waited for admission."""
+    enqueue_time = request.scheduling_enqueue_time
+    if enqueue_time is None:
+        enqueue_time = now
+    return max(0.0, now - enqueue_time)
+
+
+def drain_priority_ordered(
+    sync_queue,
+    aging_interval: float,
+    clock: Callable[[], float] = time.monotonic,
+) -> list[InferenceRequest]:
+    """Drain a finite queue snapshot and return requests in admission order."""
+    requests = []
+    for _ in range(sync_queue.qsize()):
+        try:
+            requests.append(sync_queue.get_nowait())
+        except queue.Empty:
+            break
+
+    now = clock()
+    requests.sort(
+        key=lambda request: (
+            -effective_priority(request, now, aging_interval),
+            request.scheduling_sequence
+            if request.scheduling_sequence is not None
+            else float("inf"),
+        )
+    )
+    return requests

python/infinilm/llm/request.py

@@ -21,6 +21,22 @@
 
 logger = logging.getLogger(__name__)
 
+MIN_REQUEST_PRIORITY = 0
+MAX_REQUEST_PRIORITY = 10
+
+
+def validate_request_priority(priority: int) -> int:
+    """Validate and return a request priority."""
+    error = (
+        f"`priority` must be an integer between {MIN_REQUEST_PRIORITY} "
+        f"and {MAX_REQUEST_PRIORITY}."
+    )
+    if isinstance(priority, bool) or not isinstance(priority, int):
+        raise ValueError(error)
+    if not MIN_REQUEST_PRIORITY <= priority <= MAX_REQUEST_PRIORITY:
+        raise ValueError(error)
+    return priority
+
 
 class _SequenceView(Sequence):
     """Live read-only view over an internally mutable list."""
@@ -158,6 +174,7 @@ def __init__(
         request_data: Optional[dict] = None,
         *,
         has_multimodal_inputs: bool = False,
+        priority: int = MIN_REQUEST_PRIORITY,
     ):
         self.arrival_time: float = arrival_time or time.time()
         self.finished_time: Optional[float] = None
@@ -174,7 +191,9 @@ def __init__(
         self.has_multimodal_inputs: bool = has_multimodal_inputs or bool(
             mm_token_index_mappings
         )
-        self.priority: int = 0
+        self.priority = validate_request_priority(priority)
+        self.scheduling_enqueue_time: Optional[float] = None
+        self.scheduling_sequence: Optional[int] = None
 
         # Sampling & stopping criteria
         self.sampling_params: SamplingParams = sampling_params or SamplingParams()