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

Improve dp attention port assignment scheme (#5889)

Browse Source 
Co-authored-by: Cheng Wan <cwan@x.ai>
This commit is contained in:
Yongtong Wu
2025-10-13 08:55:59 +08:00
committed by GitHub
parent 1bdd010291
commit a20e7df8d0
4 changed files with 175 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -21,7 +21,7 @@ import threading
import time
from collections import deque
from enum import Enum, auto
from typing import List
from typing import List, Optional
import psutil
import setproctitle
@@ -36,7 +36,11 @@ from sglang.srt.managers.io_struct import (
)
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.scheduler import run_scheduler_process
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.server_args import (
DP_ATTENTION_HANDSHAKE_PORT_DELTA,
PortArgs,
ServerArgs,
)
from sglang.srt.utils import (
bind_port,
configure_logger,
@@ -140,22 +144,12 @@ class DataParallelController:
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)
self.launch_dp_attention_schedulers(server_args, port_args)
self.control_message_step = server_args.tp_size
else:
dp_port_args = self.launch_dp_schedulers(server_args, port_args)
self.launch_dp_schedulers(server_args, port_args)
self.control_message_step = 1
# Only node rank 0 runs the real data parallel controller that dispatches the requests.
if server_args.node_rank == 0:
for dp_rank in range(server_args.dp_size):
self.workers[dp_rank] = get_zmq_socket(
self.context,
zmq.PUSH,
dp_port_args[dp_rank].scheduler_input_ipc_name,
True,
)
self.max_req_input_len = None
self.init_dispatcher()
@@ -188,13 +182,11 @@ class DataParallelController:
threads = []
sockets = []
dp_port_args = []
ready_events = []
for dp_rank in range(server_args.dp_size):
tmp_port_args = PortArgs.init_new(server_args)
tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
dp_port_args.append(tmp_port_args)
# This port is checked free in PortArgs.init_new.
# We hold it first so that the next dp worker gets a different port
@@ -213,6 +205,14 @@ class DataParallelController:
server_args.tp_size * server_args.pp_size * server_args.gpu_id_step
)
if server_args.node_rank == 0:
self.workers[dp_rank] = get_zmq_socket(
self.context,
zmq.PUSH,
tmp_port_args.scheduler_input_ipc_name,
True,
)
# Free all sockets before starting the threads to launch TP workers
for sock in sockets:
sock.close()
@@ -223,8 +223,6 @@ class DataParallelController:
for event in ready_events:
event.wait()
return dp_port_args
def launch_tensor_parallel_group_thread(
self,
server_args: ServerArgs,
@@ -241,19 +239,115 @@ class DataParallelController:
while True:
time.sleep(30 * 24 * 3600)
def launch_dp_attention_schedulers(self, server_args, port_args):
self.launch_tensor_parallel_group(server_args, port_args, 0, None)
dp_port_args = []
for dp_rank in range(server_args.dp_size):
dp_port_args.append(PortArgs.init_new(server_args, dp_rank))
return dp_port_args
def _broadcast_worker_ports(
self, server_args: ServerArgs, worker_ports: Optional[List[int]] = None
) -> List[int]:
"""Broadcast worker ports from node 0 to all other nodes.
Node 0 acts as the server, waiting for all other nodes to connect and
sending them the pre-allocated worker ports. Other nodes act as clients,
connecting to node 0 to receive their copy of the worker ports.
Args:
server_args: Server arguments containing node configuration.
worker_ports: Pre-allocated worker ports to broadcast.
Returns:
List of worker ports (same on all nodes after broadcast).
"""
# Determine the endpoint for inter-node communication
if server_args.dist_init_addr is None:
endpoint = f"tcp://127.0.0.1:{server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA}"
else:
endpoint = f"tcp://{server_args.dist_init_addr}"
if server_args.node_rank == 0:
# Node 0: Broadcast worker ports to all other nodes
return self._broadcast_ports_as_server(
endpoint, server_args.nnodes - 1, worker_ports
)
else:
# Other nodes: Receive worker ports from node 0
return self._receive_ports_as_client(endpoint, server_args.node_rank)
def _broadcast_ports_as_server(
self, endpoint: str, expected_clients: int, worker_ports: List[int]
) -> List[int]:
"""Broadcast worker ports to all client nodes."""
logger.debug(f"Broadcasting worker ports to {expected_clients} client nodes")
logger.debug(f"Worker ports: {worker_ports}")
rep_socket = get_zmq_socket(self.context, zmq.REP, endpoint, True)
try:
connected_clients = 0
while connected_clients < expected_clients:
# Wait for client handshake
client_rank = rep_socket.recv().decode()
logger.debug(f"Received handshake from node {client_rank}")
# Send worker ports to client
rep_socket.send_pyobj(worker_ports)
connected_clients += 1
logger.debug(
f"Sent worker ports to {connected_clients}/{expected_clients} nodes"
)
logger.debug("Worker port broadcast completed")
return worker_ports
finally:
rep_socket.close()
def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]:
"""Receive worker ports from the server node."""
logger.debug(f"Connecting to node 0 to receive worker ports")
req_socket = get_zmq_socket(self.context, zmq.REQ, endpoint, False)
req_socket.setsockopt(zmq.RCVTIMEO, 60 * 1000) # 1 minute timeout
req_socket.setsockopt(zmq.SNDTIMEO, 60 * 1000)
try:
# Send handshake with our node rank
req_socket.send(str(node_rank).encode())
# Receive worker ports
worker_ports = req_socket.recv_pyobj()
logger.debug(f"Received {len(worker_ports)} worker ports from node 0")
return worker_ports
except zmq.Again:
logger.error("Timeout waiting for worker ports from node 0")
raise RuntimeError(
"Failed to receive worker ports from node 0 within timeout"
)
finally:
req_socket.close()
def launch_dp_attention_schedulers(
self, server_args: ServerArgs, port_args: PortArgs
):
# Pre-allocate worker ports on node 0 to avoid conflicts
worker_ports = []
if server_args.node_rank == 0:
for dp_rank in range(server_args.dp_size):
port_and_socket = get_zmq_socket(self.context, zmq.PUSH)
worker_ports.append(port_and_socket[0])
self.workers[dp_rank] = port_and_socket[1]
logger.debug(f"Assigned port {port_and_socket[0]} to worker {dp_rank}")
broadcasted_ports = self._broadcast_worker_ports(
server_args, worker_ports if worker_ports else None
)
self.launch_tensor_parallel_group(
server_args, port_args, 0, None, broadcasted_ports
)
def launch_tensor_parallel_group(
self,
server_args: ServerArgs,
port_args: PortArgs,
base_gpu_id: int,
dp_rank: int,
dp_rank: Optional[int],
worker_ports: Optional[List[int]] = None,
):
if not server_args.enable_dp_attention:
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
@@ -290,7 +384,9 @@ class DataParallelController:
server_args.dp_size,
)
# compute zmq ports for this dp rank
rank_port_args = PortArgs.init_new(server_args, dp_rank)
rank_port_args = PortArgs.init_new(
server_args, dp_rank, worker_ports
)
# Data parallelism reuses the tensor parallelism group,
# so all dp ranks should use the same nccl port.
rank_port_args.nccl_port = port_args.nccl_port