feat: Priority-based scheduling optimization (including default priority, preemption toggle, priority-based metrics, etc.)

python/sglang/srt/managers/scheduler.py

@@ -19,7 +19,7 @@
 import signal
 import sys
 import time
-from collections import deque
+from collections import defaultdict, deque
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, Deque, Dict, List, Optional, Tuple, Union
@@ -786,7 +786,10 @@ def init_schedule_policy(self):
                 token_usage_low_watermark=self.server_args.prefill_delayer_token_usage_low_watermark,
             )
         # Enable preemption for priority scheduling.
-        self.try_preemption = self.enable_priority_scheduling
+        self.try_preemption = (
+            self.enable_priority_scheduling
+            and self.server_args.enable_try_preemption_by_priority
+        )
         self.init_new_token_ratio = min(
             envs.SGLANG_INIT_NEW_TOKEN_RATIO.get()
             * self.server_args.schedule_conservativeness,
@@ -1928,6 +1931,10 @@ def _get_new_batch_prefill_raw(
             return None
 
         running_bs = len(self.running_batch.reqs)
+        num_running_reqs_by_priority = defaultdict(int)
+        if self.enable_priority_scheduling:
+            for req in self.running_batch.reqs:
+                num_running_reqs_by_priority[req.priority] += 1
         # Ignore the check if self.chunked_req is not None.
         # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0,
         # as the space for the chunked requests has just been released.
@@ -2014,6 +2021,10 @@ def _get_new_batch_prefill_raw(
                         continue
 
             running_bs = len(self.running_batch.reqs)
+            if self.enable_priority_scheduling:
+                num_running_reqs_by_priority.clear()
+                for running_req in self.running_batch.reqs:
+                    num_running_reqs_by_priority[running_req.priority] += 1
             if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
                 self.running_batch.batch_is_full = True
             if self.disaggregation_mode == DisaggregationMode.PREFILL:
@@ -2100,6 +2111,7 @@ def _get_new_batch_prefill_raw(
                 can_run_list,
                 running_bs=len(self.running_batch.reqs),
                 running_bs_offline_batch=0,
+                num_running_reqs_by_priority=num_running_reqs_by_priority,
             )
 
         # Record metrics

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -107,6 +107,8 @@ def init_metrics(
                 "pp_rank": pp_rank,
                 "moe_ep_rank": self.moe_ep_rank,
             }
+            if self.enable_priority_scheduling:
+                labels["priority"] = ""
             if dp_rank is not None:
                 labels["dp_rank"] = dp_rank
             self.metrics_collector = SchedulerMetricsCollector(
@@ -150,6 +152,7 @@ def log_prefill_stats(
         can_run_list: List[Req],
         running_bs: int,
         running_bs_offline_batch: int,
+        num_running_reqs_by_priority: dict[int, int] = None,
     ):
         gap_latency = time.perf_counter() - self.last_prefill_stats_tic
         self.last_prefill_stats_tic = time.perf_counter()
@@ -229,6 +232,7 @@ def log_prefill_stats(
             )
 
             self.stats.num_running_reqs = running_bs
+            self.stats.num_running_reqs_by_priority = num_running_reqs_by_priority
             self.stats.num_running_reqs_offline_batch = running_bs_offline_batch
             self.stats.num_used_tokens = num_used
             self.stats.token_usage = token_usage
@@ -237,6 +241,11 @@ def log_prefill_stats(
             if self.is_hybrid_ssm:
                 self.stats.mamba_usage = mamba_usage
             self.stats.num_queue_reqs = len(self.waiting_queue)
+            if self.enable_priority_scheduling:
+                num_queue_reqs_by_priority: dict[int, int] = defaultdict(int)
+                for req in self.waiting_queue:
+                    num_queue_reqs_by_priority[req.priority] += 1
+                self.stats.num_queue_reqs_by_priority = num_queue_reqs_by_priority
             self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
             self.stats.cache_hit_rate = cache_hit_rate
 
@@ -255,6 +264,23 @@ def log_prefill_stats(
                 self.stats.num_prefill_inflight_queue_reqs = len(
                     self.disagg_prefill_inflight_queue
                 )
+                if self.enable_priority_scheduling:
+                    num_prefill_prealloc_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_prefill_bootstrap_queue.queue:
+                        num_prefill_prealloc_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_prefill_prealloc_queue_reqs_by_priority = (
+                        num_prefill_prealloc_queue_reqs_by_priority
+                    )
+                    num_prefill_inflight_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_prefill_inflight_queue:
+                        num_prefill_inflight_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_prefill_inflight_queue_reqs_by_priority = (
+                        num_prefill_inflight_queue_reqs_by_priority
+                    )
                 self.stats.kv_transfer_speed_gb_s = self.kv_transfer_speed_gb_s
                 self.stats.kv_transfer_latency_ms = self.kv_transfer_latency_ms
                 self.stats.kv_transfer_bootstrap_ms = self.kv_transfer_bootstrap_ms
@@ -267,6 +293,23 @@ def log_prefill_stats(
                 self.stats.num_decode_transfer_queue_reqs = len(
                     self.disagg_decode_transfer_queue.queue
                 )
+                if self.enable_priority_scheduling:
+                    num_decode_prealloc_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_decode_prealloc_queue.queue:
+                        num_decode_prealloc_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_decode_prealloc_queue_reqs_by_priority = (
+                        num_decode_prealloc_queue_reqs_by_priority
+                    )
+                    num_decode_transfer_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_decode_transfer_queue.queue:
+                        num_decode_transfer_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_decode_transfer_queue_reqs_by_priority = (
+                        num_decode_transfer_queue_reqs_by_priority
+                    )
 
             # Others
             self.calculate_utilization()
@@ -299,6 +342,10 @@ def log_decode_stats(
 
         self.num_generated_tokens = 0
         num_running_reqs = len(batch.reqs)
+        num_running_reqs_by_priority: dict[int, int] = defaultdict(int)
+        if self.enable_priority_scheduling:
+            for req in batch.reqs:
+                num_running_reqs_by_priority[req.priority] += 1
         num_running_reqs_offline_batch = 0
 
         # TODO: generalize this for various memory pools
@@ -400,6 +447,7 @@ def log_decode_stats(
         if self.enable_metrics:
             # Basics
             self.stats.num_running_reqs = num_running_reqs
+            self.stats.num_running_reqs_by_priority = num_running_reqs_by_priority
             self.stats.num_running_reqs_offline_batch = num_running_reqs_offline_batch
             self.stats.num_used_tokens = num_used
             self.stats.token_usage = token_usage
@@ -410,6 +458,11 @@ def log_decode_stats(
             self.stats.decode_sum_seq_lens = batch.seq_lens_cpu.sum().item()
             self.stats.gen_throughput = self.last_gen_throughput
             self.stats.num_queue_reqs = len(self.waiting_queue)
+            if self.enable_priority_scheduling:
+                num_queue_reqs_by_priority: dict[int, int] = defaultdict(int)
+                for req in self.waiting_queue:
+                    num_queue_reqs_by_priority[req.priority] += 1
+                self.stats.num_queue_reqs_by_priority = num_queue_reqs_by_priority
             self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
             self.stats.cache_hit_rate = cache_hit_rate
 
@@ -432,14 +485,47 @@ def log_decode_stats(
                 self.stats.num_prefill_inflight_queue_reqs = len(
                     self.disagg_prefill_inflight_queue
                 )
+                if self.enable_priority_scheduling:
+                    num_prefill_prealloc_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_prefill_bootstrap_queue.queue:
+                        num_prefill_prealloc_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_prefill_prealloc_queue_reqs_by_priority = (
+                        num_prefill_prealloc_queue_reqs_by_priority
+                    )
+                    num_prefill_inflight_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_prefill_inflight_queue:
+                        num_prefill_inflight_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_prefill_inflight_queue_reqs_by_priority = (
+                        num_prefill_inflight_queue_reqs_by_priority
+                    )
             elif self.disaggregation_mode == DisaggregationMode.DECODE:
                 self.stats.num_decode_prealloc_queue_reqs = len(
                     self.disagg_decode_prealloc_queue.queue
                 )
                 self.stats.num_decode_transfer_queue_reqs = len(
                     self.disagg_decode_transfer_queue.queue
                 )
-
+                if self.enable_priority_scheduling:
+                    num_decode_prealloc_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_decode_prealloc_queue.queue:
+                        num_decode_prealloc_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_decode_prealloc_queue_reqs_by_priority = (
+                        num_decode_prealloc_queue_reqs_by_priority
+                    )
+                    num_decode_transfer_queue_reqs_by_priority: dict[int, int] = (
+                        defaultdict(int)
+                    )
+                    for req in self.disagg_decode_transfer_queue.queue:
+                        num_decode_transfer_queue_reqs_by_priority[req.priority] += 1
+                    self.stats.num_decode_transfer_queue_reqs_by_priority = (
+                        num_decode_transfer_queue_reqs_by_priority
+                    )
             running_routing_keys = [r.routing_key for r in batch.reqs]
             waiting_routing_keys = [r.routing_key for r in self.waiting_queue]
             (

python/sglang/srt/managers/tokenizer_manager.py

@@ -245,6 +245,8 @@ def init_model_config(self):
         self.context_len = self.model_config.context_len
         self.image_token_id = self.model_config.image_token_id
         self.max_req_input_len = None  # Will be set later in engine.py
+        self.enable_priority_scheduling = server_args.enable_priority_scheduling
+        self.default_priority_value = server_args.default_priority_value
         speculative_algorithm = SpeculativeAlgorithm.from_string(
             server_args.speculative_algorithm
         )
@@ -432,6 +434,8 @@ def init_metric_collector_watchdog(self):
                 "model_name": self.server_args.served_model_name,
                 # TODO: Add lora name/path in the future,
             }
+            if self.enable_priority_scheduling:
+                labels["priority"] = ""
             if self.server_args.tokenizer_metrics_allowed_custom_labels:
                 for label in self.server_args.tokenizer_metrics_allowed_custom_labels:
                     labels[label] = ""
@@ -493,6 +497,7 @@ async def generate_request(
 
         # Normalize the request
         obj.normalize_batch_and_arguments()
+        self._set_default_priority(obj)
         if self.enable_trace:
             self._trace_request_start(obj, created_time, request)
         if self.server_args.language_only:
@@ -1934,6 +1939,10 @@ def collect_metrics(self, state: ReqState, recv_obj: BatchStrOutput, i: int):
             if custom_labels
             else self.metrics_collector.labels
         )
+        if self.enable_priority_scheduling:
+            priority = getattr(state.obj, "priority", None)
+            if priority is not None:
+                labels["priority"] = priority
         if (
             state.first_token_time == 0.0
             and self.disaggregation_mode != DisaggregationMode.PREFILL
@@ -2283,6 +2292,15 @@ def _handle_epd_disaggregation_encode_request(
         ):
             self.mm_receiver.send_encode_request(obj)
 
+    def _set_default_priority(self, obj: Union[GenerateReqInput, EmbeddingReqInput]):
+        """Set the default priority value."""
+        if (
+            self.enable_priority_scheduling
+            and obj.priority is None
+            and self.default_priority_value is not None
+        ):
+            obj.priority = self.default_priority_value
+
 
 class ServerStatus(Enum):
     Up = "Up"