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

[1/N] elastic-ep: Add runtime EP scale-up (#30164)

Browse Source
This commit is contained in:
Yoray Zack
2026-07-16 15:53:44 -07:00
committed by GitHub
parent d28e35b1a1
commit 77d23a796e
34 changed files with 2549 additions and 338 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/eplb_manager.py
+16
View File
@@ -52,6 +52,8 @@ class EPLBManager:
self._server_args.eplb_rebalance_layers_per_chunk
)
self._rebalance_num_iterations = self._server_args.eplb_rebalance_num_iterations
self._rebalance_disabled_reason = None
self._rebalance_disabled_logged = False
# Otherwise, the circular buffer will contain stale data. If the case is needed, it can be implemented.
assert (
@@ -74,6 +76,11 @@ class EPLBManager:
def reset_generator(self):
self._main_generator = self._entrypoint()
def disable_rebalance(self, reason: str):
self._rebalance_disabled_reason = reason
self._rebalance_disabled_logged = False
self.reset_generator()
# can be more complex if needed
def _entrypoint(self):
while True:
@@ -83,6 +90,15 @@ class EPLBManager:
yield from self.rebalance()
def rebalance(self):
if self._rebalance_disabled_reason is not None:
if not self._rebalance_disabled_logged:
logger.debug(
"[EPLBManager] rebalance disabled: %s",
self._rebalance_disabled_reason,
)
self._rebalance_disabled_logged = True
return
logger.info("[EPLBManager] rebalance start")
enable_timing = self._rebalance_layers_per_chunk is None
python/sglang/srt/eplb/expert_distribution.py
+18 -3
View File
@@ -331,7 +331,9 @@ class _SinglePassGatherer(ABC):
return _SelectExpertsSinglePassGatherer(expert_location_metadata, rank)
elif server_args.deepep_mode == "low_latency":
return _DeepepLowLatencySinglePassGatherer(
expert_location_metadata, rank
expert_location_metadata,
rank,
elastic_ep_enabled=server_args.elastic_ep_backend is not None,
)
else:
raise NotImplementedError
@@ -574,13 +576,26 @@ class _DeepepNormalSinglePassGatherer(_LayerBasedCpuSinglePassGatherer):
class _DeepepLowLatencySinglePassGatherer(_LayerBasedGpuSinglePassGatherer):
def __init__(self, *args, **kwargs):
def __init__(self, *args, elastic_ep_enabled: bool = False, **kwargs):
super().__init__(*args, **kwargs, enable_global_physical_experts=False)
self._elastic_ep_enabled = elastic_ep_enabled
def on_deepep_dispatch_low_latency(
self, layer_idx: int, local_physical_count_of_layer: torch.Tensor
):
# Most naive implementation, can optimize later
if local_physical_count_of_layer.shape[0] != self._data.shape[1]:
if not self._elastic_ep_enabled:
self._data[layer_idx, :] += local_physical_count_of_layer
return
n = self._data.shape[1]
if local_physical_count_of_layer.shape[0] > n:
local_physical_count_of_layer = local_physical_count_of_layer[:n]
else:
local_physical_count_of_layer = torch.nn.functional.pad(
local_physical_count_of_layer,
(0, n - local_physical_count_of_layer.shape[0]),
)
self._data[layer_idx, :] += local_physical_count_of_layer
python/sglang/srt/eplb/expert_location.py
+130 -65
View File
@@ -32,6 +32,25 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _prefer_same_node_experts(server_args: ServerArgs) -> bool:
from sglang.srt.elastic_ep.elastic_ep import elastic_expanded_world_enabled
return server_args.ep_join_mode != "scale" and not elastic_expanded_world_enabled()
def _compute_elastic_expert_layout(
base_num_physical_experts: int,
initial_ep_size: int,
effective_ep_size: int,
) -> tuple[int, int]:
assert base_num_physical_experts % initial_ep_size == 0
num_local_physical_experts = base_num_physical_experts // initial_ep_size
return (
num_local_physical_experts * effective_ep_size,
num_local_physical_experts,
)
@dataclass
class ExpertLocationMetadata:
physical_to_logical_map: torch.Tensor # (layers, num_physical_experts)
@@ -39,6 +58,7 @@ class ExpertLocationMetadata:
logical_to_all_physical_map: torch.Tensor # (layers, num_logical_experts, X)
logical_to_all_physical_map_cpu: torch.Tensor # CPU copy for performance
logical_to_all_physical_map_num_valid: torch.Tensor # (layers, num_logical_experts)
ep_size: int
# (layers, num_logical_experts)
logical_to_rank_dispatch_physical_map: Optional[torch.Tensor]
@@ -62,11 +82,6 @@ class ExpertLocationMetadata:
def num_logical_experts(self) -> int:
return self.logical_to_all_physical_map.shape[1]
@property
def ep_size(self):
# TODO change when EP size != world size
return torch.distributed.get_world_size()
def __post_init__(self):
num_layers_0, num_physical_experts_0 = self.physical_to_logical_map.shape
num_layers_1, num_logical_experts_0, num_physical_experts_1 = (
@@ -96,10 +111,16 @@ class ExpertLocationMetadata:
num_layers = model_config_for_expert_location.num_layers
num_logical_experts = model_config_for_expert_location.num_logical_experts
base_num_physical_experts = common["base_num_physical_experts"]
physical_to_logical_map = (
torch.arange(0, num_physical_experts).repeat(num_layers, 1)
torch.arange(0, base_num_physical_experts).repeat(num_layers, 1)
% num_logical_experts
)
physical_to_logical_map = append_trivial_expert_slots(
physical_to_logical_map,
num_physical_experts - base_num_physical_experts,
num_logical_experts,
)
return ExpertLocationMetadata.init_by_mapping(
server_args,
@@ -125,6 +146,15 @@ class ExpertLocationMetadata:
return None
model_config_for_expert_location = common["model_config_for_expert_location"]
if common["num_physical_experts"] > common["base_num_physical_experts"]:
if physical_to_logical_map.shape[-1] == common["base_num_physical_experts"]:
physical_to_logical_map = append_trivial_expert_slots(
physical_to_logical_map,
common["num_physical_experts"]
- common["base_num_physical_experts"],
model_config_for_expert_location.num_logical_experts,
)
assert physical_to_logical_map.shape[-1] == common["num_physical_experts"]
logical_to_all_physical_map = _compute_logical_to_all_physical_map(
server_args=server_args,
physical_to_logical_map=physical_to_logical_map,
@@ -138,6 +168,7 @@ class ExpertLocationMetadata:
ep_size=common["ep_size"],
physical_to_logical_map=physical_to_logical_map,
logical_to_all_physical_map=logical_to_all_physical_map,
moe_ep_rank=moe_ep_rank,
)
@staticmethod
@@ -195,16 +226,33 @@ class ExpertLocationMetadata:
if model_config_for_expert_location is None:
return None
num_physical_experts = (
base_num_physical_experts = (
model_config_for_expert_location.num_logical_experts
+ server_args.ep_num_redundant_experts
)
ep_size = server_args.ep_size
assert num_physical_experts % ep_size == 0
num_local_physical_experts = num_physical_experts // ep_size
num_physical_experts = base_num_physical_experts
initial_ep_size = server_args.elastic_ep_initial_size
if initial_ep_size is not None:
if server_args.ep_join_mode == "scale":
ep_size = max(
ep_size,
server_args.ep_join_rank_offset + server_args.tp_size,
)
num_physical_experts, num_local_physical_experts = (
_compute_elastic_expert_layout(
base_num_physical_experts,
initial_ep_size,
ep_size,
)
)
else:
assert num_physical_experts % ep_size == 0
num_local_physical_experts = num_physical_experts // ep_size
return dict(
model_config_for_expert_location=model_config_for_expert_location,
base_num_physical_experts=base_num_physical_experts,
num_physical_experts=num_physical_experts,
num_local_physical_experts=num_local_physical_experts,
ep_size=ep_size,
@@ -216,6 +264,7 @@ class ExpertLocationMetadata:
ep_size: int,
physical_to_logical_map: torch.Tensor,
logical_to_all_physical_map: torch.Tensor,
moe_ep_rank: Optional[int] = None,
):
_, num_physical_experts = physical_to_logical_map.shape
@@ -235,14 +284,18 @@ class ExpertLocationMetadata:
logical_to_all_physical_map=logical_to_all_physical_map_padded,
logical_to_all_physical_map_cpu=logical_to_all_physical_map_padded.cpu(),
logical_to_all_physical_map_num_valid=logical_to_all_physical_map_num_valid,
ep_size=ep_size,
logical_to_rank_dispatch_physical_map=(
compute_logical_to_rank_dispatch_physical_map(
server_args=server_args,
logical_to_all_physical_map=logical_to_all_physical_map,
ep_size=ep_size,
num_physical_experts=num_physical_experts,
# TODO improve when we have real EP rank
ep_rank=torch.distributed.get_rank() % ep_size,
ep_rank=(
moe_ep_rank
if moe_ep_rank is not None
else torch.distributed.get_rank() % ep_size
),
)
if server_args.ep_dispatch_algorithm == "static"
else None
@@ -425,58 +478,57 @@ def get_global_expert_location_metadata():
return get_resources().expert_location_metadata
def set_global_expert_location_metadata(value):
def set_global_expert_location_metadata(value, allow_overwrite=False):
from sglang.srt.runtime_context import get_resources
resources = get_resources()
assert resources.expert_location_metadata is None
if not allow_overwrite:
assert resources.expert_location_metadata is None
resources.expert_location_metadata = value
def append_trivial_expert_slots(
physical_to_logical_map: torch.Tensor,
count: int,
num_logical_experts: int,
start: int = 0,
) -> torch.Tensor:
if count <= 0:
return physical_to_logical_map
new_slots = torch.arange(
start,
start + count,
dtype=physical_to_logical_map.dtype,
device=physical_to_logical_map.device,
).unsqueeze(0)
new_slots = new_slots.expand(physical_to_logical_map.shape[0], -1)
return torch.cat([physical_to_logical_map, new_slots % num_logical_experts], dim=1)
def broadcast_global_expert_location_metadata(
src_rank: int = 0, group: Optional[torch.distributed.ProcessGroup] = None
):
"""Broadcast the global ExpertLocationMetadata from src_rank to all ranks.
model_config: ModelConfig,
moe_ep_rank: int,
src_rank: int = 0,
group: Optional[torch.distributed.ProcessGroup] = None,
) -> ExpertLocationMetadata:
from sglang.srt.runtime_context import get_server_args
This is used in Elastic EP rank recovery to ensure that all ranks (including
newly recovered ones) share exactly the same expert location metadata.
Note: The caller must ensure src_rank is a healthy rank. In recovery scenarios,
this function is called after try_recover_ranks succeeds, at which point all
ranks (including src_rank=0) have recovered and are ready.
"""
server_args = get_server_args()
metadata = get_global_expert_location_metadata()
assert metadata is not None
# Ensure device tensors are contiguous before broadcasting in-place
metadata.physical_to_logical_map = metadata.physical_to_logical_map.contiguous()
metadata.logical_to_all_physical_map = (
metadata.logical_to_all_physical_map.contiguous()
torch.distributed.broadcast(
metadata.physical_to_logical_map, src=src_rank, group=group
)
metadata.logical_to_all_physical_map_num_valid = (
metadata.logical_to_all_physical_map_num_valid.contiguous()
)
if metadata.logical_to_rank_dispatch_physical_map is not None:
metadata.logical_to_rank_dispatch_physical_map = (
metadata.logical_to_rank_dispatch_physical_map.contiguous()
)
device_tensors = [
metadata = ExpertLocationMetadata.init_by_mapping(
server_args,
model_config,
metadata.physical_to_logical_map,
metadata.logical_to_all_physical_map,
metadata.logical_to_all_physical_map_num_valid,
]
if metadata.logical_to_rank_dispatch_physical_map is not None:
device_tensors.append(metadata.logical_to_rank_dispatch_physical_map)
for tensor in device_tensors:
torch.distributed.broadcast(tensor, src=src_rank, group=group)
# After broadcasting device tensors, refresh corresponding CPU copies
metadata.physical_to_logical_map_cpu = metadata.physical_to_logical_map.cpu()
metadata.logical_to_all_physical_map_cpu = (
metadata.logical_to_all_physical_map.cpu()
moe_ep_rank=moe_ep_rank,
)
set_global_expert_location_metadata(metadata, allow_overwrite=True)
return metadata
def _compute_logical_to_all_physical_map(
@@ -506,10 +558,15 @@ def _compute_logical_to_all_physical_map(
# Replace by the physical expert on local GPU or node if possible
if moe_ep_rank is not None:
num_gpus_per_node = server_args.ep_size // server_args.nnodes
num_local_gpu_physical_experts = num_physical_experts // ep_size
prefer_same_node = _prefer_same_node_experts(server_args)
num_gpus_per_node = (
server_args.ep_size // server_args.nnodes if prefer_same_node else None
)
num_local_node_physical_experts = (
num_local_gpu_physical_experts * num_gpus_per_node
if num_gpus_per_node is not None
else None
)
for layer_id in range(num_layers):
for logical_expert_id in range(num_logical_experts):
@@ -563,8 +620,15 @@ def compute_logical_to_rank_dispatch_physical_map(
logical_to_all_physical_map = logical_to_all_physical_map.cpu()
num_local_gpu_physical_experts = num_physical_experts // ep_size
num_gpus_per_node = server_args.ep_size // server_args.nnodes
num_local_node_physical_experts = num_local_gpu_physical_experts * num_gpus_per_node
prefer_same_node = _prefer_same_node_experts(server_args)
num_gpus_per_node = (
server_args.ep_size // server_args.nnodes if prefer_same_node else None
)
num_local_node_physical_experts = (
num_local_gpu_physical_experts * num_gpus_per_node
if num_gpus_per_node is not None
else None
)
num_layers, num_logical_experts, _ = logical_to_all_physical_map.shape
dtype = logical_to_all_physical_map.dtype
@@ -633,8 +697,8 @@ def _find_nearest_expert(
candidate_physical_expert_ids: List[int],
num_local_gpu_physical_experts: int,
moe_ep_rank: int,
num_gpus_per_node: int,
num_local_node_physical_experts: int,
num_gpus_per_node: Optional[int],
num_local_node_physical_experts: Optional[int],
) -> int:
# 1. If only one candidate, return it directly
if len(candidate_physical_expert_ids) == 1:
@@ -652,18 +716,19 @@ def _find_nearest_expert(
if len(same_gpu_physical_expert_ids) > 0:
return same_gpu_physical_expert_ids[0]
# 3. Otherwise, prefer same-node experts
node_rank = moe_ep_rank // num_gpus_per_node
same_node_physical_expert_ids = [
physical_expert_id
for physical_expert_id in candidate_physical_expert_ids
if _compute_node_id_of_physical_expert(
physical_expert_id, num_local_node_physical_experts
)
== node_rank
]
if len(same_node_physical_expert_ids) > 0:
return same_node_physical_expert_ids[0]
# Prefer same-node experts only when it narrows the candidate set.
if num_gpus_per_node is not None and num_local_node_physical_experts is not None:
node_rank = moe_ep_rank // num_gpus_per_node
same_node_physical_expert_ids = [
physical_expert_id
for physical_expert_id in candidate_physical_expert_ids
if _compute_node_id_of_physical_expert(
physical_expert_id, num_local_node_physical_experts
)
== node_rank
]
if 0 < len(same_node_physical_expert_ids) < len(candidate_physical_expert_ids):
return same_node_physical_expert_ids[0]
# 4. At last, leave it as -1 to indicate not found.
return -1