fix(disagg): broadcast bootstrap port across multi-node prefill ranks (#24378)

This commit is contained in:
YAMY
2026-05-14 16:39:01 +08:00
committed by GitHub
parent 2417a9da57
commit 4be25f2428
@@ -12,6 +12,7 @@ from typing import Dict, List, Optional, Set, Tuple, Union
import numpy as np import numpy as np
import numpy.typing as npt import numpy.typing as npt
import requests import requests
import torch.distributed as dist
import zmq import zmq
from aiohttp import web from aiohttp import web
@@ -25,7 +26,7 @@ from sglang.srt.disaggregation.base.conn import (
KVTransferMetric, KVTransferMetric,
) )
from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.distributed import get_pp_group from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
get_attention_cp_rank, get_attention_cp_rank,
@@ -142,6 +143,13 @@ class CommonKVManager(BaseKVManager):
and self.attn_cp_size > 1 and self.attn_cp_size > 1
and self.attn_cp_rank != 0 and self.attn_cp_rank != 0
) )
# Sync the leader's bootstrap port to every rank before
# registering: in multi-node prefill, registration targets
# `dist_init_addr` (rank 0) but each rank's local port may
# differ when the launcher auto-reserves a free port per host.
self.bootstrap_port = self._sync_bootstrap_port_across_nodes(
self.bootstrap_port
)
self.register_to_bootstrap() self.register_to_bootstrap()
self.transfer_infos = {} self.transfer_infos = {}
self.req_to_decode_prefix_len: Dict[int, int] = {} self.req_to_decode_prefix_len: Dict[int, int] = {}
@@ -335,6 +343,36 @@ class CommonKVManager(BaseKVManager):
info.required_dst_info_num = required_dst_info_num info.required_dst_info_num = required_dst_info_num
info.required_prefill_response_num = required_prefill_response_num info.required_prefill_response_num = required_prefill_response_num
def _sync_bootstrap_port_across_nodes(self, local_port: int) -> int:
"""Broadcast world-rank-0's bootstrap port to all prefill ranks.
Required for multi-node prefill when the launcher auto-reserves a
free port per host (e.g. Dynamo's
`_reserve_disaggregation_bootstrap_port`): without sync, non-leader
ranks register to `<leader_ip>:<their_local_port>`, hit
`Connection refused`, and the leader's `prefill_port_table` ends
up missing rows.
"""
if not self.dist_init_addr or self.server_args.nnodes == 1:
return local_port
if not (dist.is_available() and dist.is_initialized()):
raise RuntimeError(
"torch.distributed must be initialised before "
"CommonKVManager registers to the bootstrap server in "
"multi-node prefill mode."
)
world_group = get_world_group()
synced_port = world_group.broadcast_object(local_port, src=0)
if synced_port != local_port:
logger.info(
f"Synced disaggregation bootstrap port from leader: "
f"local={local_port} -> leader={synced_port} "
f"(world_rank={world_group.rank_in_group})"
)
return synced_port
def register_to_bootstrap(self): def register_to_bootstrap(self):
"""Register prefill server info to bootstrap server via HTTP POST.""" """Register prefill server info to bootstrap server via HTTP POST."""
if self.dist_init_addr: if self.dist_init_addr: