diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index ec9ae1d3d..ca75f5b43 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -68,6 +68,7 @@ class LoadBalanceMethod(Enum): """Load balance method.""" ROUND_ROBIN = auto() + DECODE_ROUND_ROBIN = auto() SHORTEST_QUEUE = auto() MINIMUM_TOKENS = auto() @@ -147,6 +148,7 @@ class DataParallelController: self.round_robin_counter = 0 dispatch_lookup = { LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler, + LoadBalanceMethod.DECODE_ROUND_ROBIN: self.decode_round_robin_scheduler, LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler, LoadBalanceMethod.MINIMUM_TOKENS: self.minimum_tokens_scheduler, } @@ -483,6 +485,18 @@ class DataParallelController: ), "req.bootstrap_room should not be None. Do not send requests directly to prefill or decode instances, but send to the router instead." self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req) + def decode_round_robin_scheduler(self, req: Req): + if self.maybe_external_dp_rank_routing(req): + return + + if self.server_args.disaggregation_mode == "decode": + self.workers[self.round_robin_counter].send_pyobj(req) + self.round_robin_counter = (self.round_robin_counter + 1) % len( + self.workers + ) + return + self.round_robin_scheduler(req) + def shortest_queue_scheduler(self, req): if self.maybe_external_dp_rank_routing(req): return diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 72cee0a09..073584044 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -2071,6 +2071,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi if ( self.server_args.dp_size == 1 or self.server_args.load_balance_method == "round_robin" + or self.server_args.load_balance_method == "decode_round_robin" ): return diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3eeabca04..b0f7ba891 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3136,6 +3136,7 @@ class ServerArgs: help="The load balancing strategy for data parallelism.", choices=[ "round_robin", + "decode_round_robin", "shortest_queue", "minimum_tokens", ],