Expose more arguments to control the scheduling policy (#32)

python/sglang/srt/managers/router/model_rpc.py

@@ -45,6 +45,7 @@ class ModelRpcServer(rpyc.Service):
         self.tp_rank = tp_rank
         self.tp_size = server_args.tp_size
         self.schedule_heuristic = server_args.schedule_heuristic
+        self.schedule_conservativeness = server_args.schedule_conservativeness
 
         # Init model and tokenizer
         self.model_config = ModelConfig(
@@ -248,7 +249,9 @@ class ModelRpcServer(rpyc.Service):
         available_size = (
             self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
         )
-        new_ratio = self.scheduler.new_token_estimation_ratio()
+        new_ratio = (
+            self.scheduler.new_token_estimation_ratio() * self.schedule_conservativeness
+        )
         if self.running_batch:
             available_size -= sum(
                 [

python/sglang/srt/server_args.py

@@ -16,6 +16,7 @@ class ServerArgs:
     tp_size: int = 1
     model_mode: List[str] = ()
     schedule_heuristic: str = "lpm"
+    schedule_conservativeness: float = 1.0
     random_seed: int = 42
     stream_interval: int = 2
     disable_log_stats: bool = False
@@ -85,7 +86,7 @@ class ServerArgs:
             "--mem-fraction-static",
             type=float,
             default=ServerArgs.mem_fraction_static,
-            help="The fraction of the memory used for static allocation (model weights and KV cache memory pool)",
+            help="The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors.",
         )
         parser.add_argument(
             "--tp-size",
@@ -107,6 +108,12 @@ class ServerArgs:
             default=ServerArgs.schedule_heuristic,
             help="Schudule mode: [lpm, weight, random, fcfs]",
         )
+        parser.add_argument(
+            "--schedule-conservativeness",
+            type=float,
+            default=ServerArgs.schedule_conservativeness,
+            help="How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see out-of-memory errors.",
+        )
         parser.add_argument(
             "--random-seed",
             type=int,