chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[1/N]DP refactor: Improve dp rank scheduling in PD disaggregation mode. (#10169)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-09-09 12:27:55 +08:00
committed by GitHub
parent 2fe17735a6
commit 83d55ac51f
8 changed files with 61 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -106,7 +106,7 @@ class DataParallelController:
# Launch data parallel workers
self.scheduler_procs = []
self.workers = [None] * server_args.dp_size
self.workers: List[zmq.Socket] = [None] * server_args.dp_size
if server_args.enable_dp_attention:
dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args)
@@ -272,27 +272,34 @@ class DataParallelController:
self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
def maybe_external_dp_rank_routing(self, req: Req):
if req.data_parallel_rank is not None:
logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
self.workers[req.data_parallel_rank].send_pyobj(req)
return True
return False
def round_robin_scheduler(self, req: Req):
if self.maybe_external_dp_rank_routing(req):
return
if self.server_args.disaggregation_mode == "null":
if req.data_parallel_rank is not None:
logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
self.workers[req.data_parallel_rank].send_pyobj(req)
else:
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(
self.workers
)
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(
self.workers
)
else:
if req.data_parallel_rank is not None:
logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
self.workers[req.data_parallel_rank].send_pyobj(req)
else:
self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
def shortest_queue_scheduler(self, input_requests):
if self.maybe_external_dp_rank_routing(req):
return
raise NotImplementedError()
def minimum_tokens_scheduler(self, req):
if self.maybe_external_dp_rank_routing(req):
return
# This variable corresponds to the balance_id in TokenizedGenerateReqInput.
# We use it to to control the number of onfly tokens (requests dispatched to workers but not yet received).
def get_next_global_balance_id() -> int: