diff --git a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py index 8814f7f702..673d76bcc0 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py @@ -28,6 +28,7 @@ HealthCheckOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, + WatchLoadUpdateReq, ) from sglang.srt.observability.req_time_stats import ( APIServerReqTimeStats, @@ -679,6 +680,15 @@ async def cleanup(request_id): cleanup_tasks.append(asyncio.create_task(cleanup(rid))) + # Forward load info to DataParallelController for token-aware balancing. + # Mirrors TokenizerManager._handle_batch_output logic: when dp_size > 1, + # each scheduler piggybacks its load (num_reqs, num_tokens) on batch output. + # Without this, DPBudget stays at zero and total_tokens/total_requests + # policies degenerate to always picking rank 0. + if self.server_args.dp_size > 1 and batch_out.load is not None: + load_update = WatchLoadUpdateReq(loads=[batch_out.load]) + self.send_to_scheduler.send_pyobj(load_update) + # Execute all queue.put() operations in parallel if put_tasks: await asyncio.gather(*put_tasks, return_exceptions=True)