support dp attn with client lb (#33105)
Co-authored-by: Kan Wu <wukanustc@gmail.com>
This commit is contained in:
@@ -76,6 +76,13 @@ class RustServer:
|
|||||||
"drop --preferred-sampling-params and send those values per request."
|
"drop --preferred-sampling-params and send those values per request."
|
||||||
)
|
)
|
||||||
http_addr = f"{server_args.host}:{server_args.port}"
|
http_addr = f"{server_args.host}:{server_args.port}"
|
||||||
|
|
||||||
|
# Per-DP-rank HTTP port with client load balancing. `None` when DP is off,
|
||||||
|
# so the rank is not conflated with rank 0 of a one-rank group.
|
||||||
|
dp_rank = scheduler.ps.attn_dp_rank if scheduler.ps.dp_size > 1 else None
|
||||||
|
if dp_rank is not None:
|
||||||
|
http_addr = f"{server_args.host}:{server_args.port + dp_rank}"
|
||||||
|
|
||||||
launch_cores, server_cores = cls._partition_cores()
|
launch_cores, server_cores = cls._partition_cores()
|
||||||
|
|
||||||
server = Server(
|
server = Server(
|
||||||
@@ -93,9 +100,15 @@ class RustServer:
|
|||||||
except OSError as e:
|
except OSError as e:
|
||||||
logger.warning("rust server: cannot pin scheduler launch thread: %s", e)
|
logger.warning("rust server: cannot pin scheduler launch thread: %s", e)
|
||||||
|
|
||||||
|
# Under DP every rank runs its own server on its own port, so the rank is
|
||||||
|
# what tells two otherwise identical startup lines apart.
|
||||||
|
dp_note = (
|
||||||
|
"" if dp_rank is None else f" (DP rank {dp_rank}/{scheduler.ps.dp_size})"
|
||||||
|
)
|
||||||
logger.info(
|
logger.info(
|
||||||
"SGLANG_RUST_SERVER enabled, Rust server listen on %s",
|
"SGLANG_RUST_SERVER enabled, Rust server listen on %s%s",
|
||||||
http_addr,
|
http_addr,
|
||||||
|
dp_note,
|
||||||
)
|
)
|
||||||
|
|
||||||
return cls(server)
|
return cls(server)
|
||||||
|
|||||||
@@ -9332,6 +9332,10 @@ class PortArgs:
|
|||||||
if dp_rank is None:
|
if dp_rank is None:
|
||||||
# TokenizerManager to DataParallelController
|
# TokenizerManager to DataParallelController
|
||||||
scheduler_input_port = port_base + 4
|
scheduler_input_port = port_base + 4
|
||||||
|
elif is_rust_server:
|
||||||
|
# Rust server path (SGLANG_RUST_SERVER + dp attention): there is no
|
||||||
|
# DataParallelController allocating worker ports.
|
||||||
|
scheduler_input_port = port_base + 6 + dp_rank
|
||||||
else:
|
else:
|
||||||
assert worker_ports is not None
|
assert worker_ports is not None
|
||||||
scheduler_input_port = worker_ports[dp_rank]
|
scheduler_input_port = worker_ports[dp_rank]
|
||||||
|
|||||||
Reference in New Issue
Block a user