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

Improve EPLB dispatch handling and diagnostics (#30646)

Browse Source 
Co-authored-by: Yonghao Zhuang <yhzhuang@meta.com>
This commit is contained in:
Lianmin Zheng
2026-07-10 10:40:19 -07:00
committed by GitHub
co-authored by Yonghao Zhuang
parent 789bc3995c
commit 3dc93a12ca
8 changed files with 222 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/eplb_manager.py
+68 -3
View File
@@ -4,8 +4,14 @@ from typing import TYPE_CHECKING, List
import torch.cuda
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ExpertLocationMetadata
from sglang.srt.eplb.expert_location import (
ExpertLocationMetadata,
format_expert_location_layout,
format_expert_location_layout_diff,
get_global_expert_location_metadata,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
@@ -78,14 +84,23 @@ class EPLBManager:
)
update_layer_ids_chunks = self._compute_update_layer_ids_chunks()
for chunk_index, update_layer_ids in enumerate(update_layer_ids_chunks):
all_update_layer_ids = [
layer_id for chunk in update_layer_ids_chunks for layer_id in chunk
]
self._log_rebalance_layout_before_update(
expert_location_metadata,
update_layer_ids=all_update_layer_ids,
)
for chunk_layer_ids in update_layer_ids_chunks:
if len(update_layer_ids_chunks) > 1:
yield
self._model_runner.update_expert_location(
expert_location_metadata,
update_layer_ids=update_layer_ids,
update_layer_ids=chunk_layer_ids,
)
self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
msg = f"[EPLBManager] rebalance end"
if enable_timing:
torch.get_device_module().synchronize()
@@ -115,6 +130,56 @@ class EPLBManager:
chunk_size = self._rebalance_layers_per_chunk or 1000000
return list(_chunk_list(all_layer_ids, chunk_size=chunk_size))
def _should_log_expert_location_metadata(self) -> bool:
return (
self._model_runner.tp_rank == 0
and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get()
)
def _log_rebalance_layout_before_update(
self,
new_expert_location_metadata: ExpertLocationMetadata,
update_layer_ids: List[int],
):
if not self._should_log_expert_location_metadata():
return
old_expert_location_metadata = get_global_expert_location_metadata()
logger.info(
"[EPLBManager] rebalance layout before:\n%s",
format_expert_location_layout(
old_expert_location_metadata,
layer_ids=update_layer_ids,
),
)
logger.info(
"[EPLBManager] rebalance layout target:\n%s",
format_expert_location_layout(
new_expert_location_metadata,
layer_ids=update_layer_ids,
),
)
logger.info(
"[EPLBManager] rebalance layout diff:\n%s",
format_expert_location_layout_diff(
old_expert_location_metadata,
new_expert_location_metadata,
layer_ids=update_layer_ids,
),
)
def _log_rebalance_layout_after_update(self, update_layer_ids: List[int]):
if not self._should_log_expert_location_metadata():
return
logger.info(
"[EPLBManager] rebalance layout after:\n%s",
format_expert_location_layout(
get_global_expert_location_metadata(),
layer_ids=update_layer_ids,
),
)
def _chunk_list(items: List, chunk_size):
for start_index in range(0, len(items), chunk_size):
python/sglang/srt/eplb/expert_distribution.py
+9 -3
View File
@@ -326,7 +326,7 @@ class _SinglePassGatherer(ABC):
else:
raise NotImplementedError
if server_args.moe_a2a_backend != "none":
if server_args.moe_a2a_backend == "deepep":
if server_args.deepep_mode == "normal":
return _SelectExpertsSinglePassGatherer(expert_location_metadata, rank)
elif server_args.deepep_mode == "low_latency":
@@ -336,6 +336,8 @@ class _SinglePassGatherer(ABC):
else:
raise NotImplementedError
# Non-DeepEP a2a backends (flashinfer / nixl / mooncake / megamoe) and
# no-a2a path dispatch through the standard topk select_experts.
return _SelectExpertsSinglePassGatherer(expert_location_metadata, rank)
def __init__(self, expert_location_metadata: ExpertLocationMetadata, rank: int):
@@ -736,11 +738,15 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
utilization_rate_gpu = torch.mean(
compute_utilization_rate(gpu_physical_count)
)
should_track_history = not math.isclose(
self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0
)
if envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get():
print(f"hi {self._rank=} {utilization_rate_gpu=}")
outputs["metrics"] = ExpertDistributionMetrics(
eplb_balancedness=utilization_rate_gpu,
)
if should_track_history:
self._history.append(utilization_rate_gpu.item())
else:
# TODO maybe refactor this part to also avoid a `.item()` gpu->cpu sync
utilization_rate_cpu = utilization_rate_gpu.item()
@@ -794,7 +800,7 @@ class _DequeCollection:
d.clear()
def mean(self) -> Dict[int, float]:
return {d.maxlen: sum(d) / len(d) for d in self._dequeues}
return {d.maxlen: sum(d) / len(d) for d in self._dequeues if len(d) > 0}
class _DetailAccumulator(_UtilizationRateAccumulatorMixin):
python/sglang/srt/eplb/expert_location.py
+104 -1
View File
@@ -19,7 +19,7 @@ import logging
import random
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, List, Optional
from typing import TYPE_CHECKING, Iterable, List, Optional
import torch
import torch.distributed
@@ -305,6 +305,109 @@ class ExpertLocationMetadata:
]
def format_expert_location_layout(
metadata: Optional[ExpertLocationMetadata],
layer_ids: Optional[Iterable[int]] = None,
) -> str:
if metadata is None:
return "<none>"
return format_physical_to_logical_map(
metadata.physical_to_logical_map_cpu,
ep_size=metadata.ep_size,
layer_ids=layer_ids,
)
def format_expert_location_layout_diff(
old_metadata: Optional[ExpertLocationMetadata],
new_metadata: Optional[ExpertLocationMetadata],
layer_ids: Optional[Iterable[int]] = None,
) -> str:
if old_metadata is None or new_metadata is None:
return "<none>"
old_map = old_metadata.physical_to_logical_map_cpu
new_map = new_metadata.physical_to_logical_map_cpu
if old_map.shape != new_map.shape:
return f"shape_changed old_shape={tuple(old_map.shape)} new_shape={tuple(new_map.shape)}"
layer_ids = _normalize_layer_ids(layer_ids, num_layers=old_map.shape[0])
num_physical_experts = old_map.shape[1]
changed_by_layer = []
for layer_id in layer_ids:
num_changed = torch.count_nonzero(old_map[layer_id] != new_map[layer_id]).item()
if num_changed > 0:
changed_by_layer.append((layer_id, num_changed))
total_changed = sum(num_changed for _, num_changed in changed_by_layer)
total_slots = len(layer_ids) * num_physical_experts
lines = [f"changed_physical_slots={total_changed}/{total_slots}"]
if not changed_by_layer:
lines.append("changed_layers=[]")
return "\n".join(lines)
for layer_id, num_changed in changed_by_layer:
lines.append(f"layer={layer_id}: changed={num_changed}/{num_physical_experts}")
return "\n".join(lines)
def format_physical_to_logical_map(
physical_to_logical_map: torch.Tensor,
ep_size: int,
layer_ids: Optional[Iterable[int]] = None,
) -> str:
physical_to_logical_map = physical_to_logical_map.cpu()
if physical_to_logical_map.numel() == 0:
return "<empty>"
layer_ids = _normalize_layer_ids(
layer_ids, num_layers=physical_to_logical_map.shape[0]
)
num_physical_experts = physical_to_logical_map.shape[1]
num_local_physical_experts, remainder = divmod(num_physical_experts, ep_size)
lines = [
"physical_to_logical_map "
f"num_layers={physical_to_logical_map.shape[0]} "
f"num_physical_experts={num_physical_experts} "
f"ep_size={ep_size}"
]
for layer_id in layer_ids:
row = physical_to_logical_map[layer_id].tolist()
if remainder != 0:
lines.append(
f"layer={layer_id}: "
f"physical={json.dumps(row, separators=(',', ':'))}"
)
continue
rank_chunks = []
for ep_rank in range(ep_size):
start = ep_rank * num_local_physical_experts
end = start + num_local_physical_experts
rank_chunks.append(
f"ep{ep_rank}={json.dumps(row[start:end], separators=(',', ':'))}"
)
lines.append(f"layer={layer_id}: " + " ".join(rank_chunks))
return "\n".join(lines)
def _normalize_layer_ids(
layer_ids: Optional[Iterable[int]],
num_layers: int,
) -> List[int]:
if layer_ids is None:
return list(range(num_layers))
normalized_layer_ids = [int(layer_id) for layer_id in layer_ids]
for layer_id in normalized_layer_ids:
assert 0 <= layer_id < num_layers, f"{layer_id=} {num_layers=}"
return normalized_layer_ids
def get_global_expert_location_metadata():
from sglang.srt.runtime_context import get_resources
python/sglang/srt/eplb/expert_location_dispatch.py
+2 -5
View File
@@ -19,9 +19,6 @@ import torch
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import is_hip
_is_hip = is_hip()
@dataclass
@@ -102,7 +99,7 @@ def _topk_ids_logical_to_physical_static(
topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
) -> torch.Tensor:
physical_topk_ids = info.partial_logical_to_rank_dispatch_physical_map[topk_ids]
if _is_hip:
if physical_topk_ids.dtype != topk_ids.dtype:
physical_topk_ids = physical_topk_ids.to(topk_ids.dtype)
return physical_topk_ids
@@ -120,7 +117,7 @@ def _topk_ids_logical_to_physical_dynamic(
% info.partial_logical_to_all_physical_map_num_valid[topk_ids]
)
topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index]
if _is_hip:
if topk_ids.dtype != original_dtype:
topk_ids = topk_ids.to(original_dtype)
topk_ids = topk_ids.view(topk_ids_original_shape)
python/sglang/srt/layers/moe/topk.py
-1
View File
@@ -2116,7 +2116,6 @@ def select_experts(
assert (
num_token_non_padded is None
), "num_token_non_padded is not yet supported in custom_routing_function"
assert expert_location_dispatch_info is None
assert not apply_routed_scaling_factor_on_output, "Not implemented"
topk_weights, topk_ids = custom_routing_function(
hidden_states=hidden_states,
python/sglang/srt/model_executor/model_runner.py
+5 -1
View File
@@ -101,6 +101,7 @@ from sglang.srt.eplb.expert_location import (
ExpertLocationMetadata,
broadcast_global_expert_location_metadata,
compute_initial_expert_location_metadata,
format_expert_location_layout,
get_global_expert_location_metadata,
set_global_expert_location_metadata,
)
@@ -670,7 +671,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
logger.info(
f"Initial expert_location_metadata: {get_global_expert_location_metadata()}"
"Initial expert_location_metadata:\n%s",
format_expert_location_layout(
get_global_expert_location_metadata()
),
)
set_global_expert_distribution_recorder(
python/sglang/srt/utils/common.py
+23
View File
@@ -1870,7 +1870,27 @@ def suppress_other_loggers():
logging.getLogger("vllm.config").setLevel(logging.ERROR)
_KERNEL_VERSION_CHECK_PACKAGES = frozenset(
{
"flashinfer-python",
"flashinfer_python",
"sglang-kernel",
"sglang_kernel",
}
)
def _should_skip_kernel_pkg_version_check(pkg: str) -> bool:
return (
pkg in _KERNEL_VERSION_CHECK_PACKAGES
and envs.SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK.get()
)
def assert_pkg_version(pkg: str, min_version: str, message: str):
if _should_skip_kernel_pkg_version_check(pkg):
return
try:
installed_version = version(pkg)
if pkg_version.parse(installed_version) < pkg_version.parse(min_version):
@@ -1896,6 +1916,9 @@ def check_pkg_version_at_least(pkg: str, min_version: str) -> bool:
Returns:
True if package is installed and version >= min_version, False otherwise
"""
if _should_skip_kernel_pkg_version_check(pkg):
return True
try:
installed_version = version(pkg)
return pkg_version.parse(installed_version) >= pkg_version.parse(min_version)