diff --git a/python/sglang/srt/model_executor/expert_location_updater.py b/python/sglang/srt/model_executor/expert_location_updater.py index 75fd7c2d7..5dce3d9a3 100644 --- a/python/sglang/srt/model_executor/expert_location_updater.py +++ b/python/sglang/srt/model_executor/expert_location_updater.py @@ -30,6 +30,9 @@ from sglang.srt.utils import get_bool_env_var logger = logging.getLogger(__name__) +_LOG_INPUT = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_INPUT") + + class ExpertLocationUpdater: def __init__(self): self._first_execution = True @@ -175,6 +178,19 @@ def update_expert_weights_single_layer( assert isinstance(old_physical_to_logical_map, list) assert isinstance(new_physical_to_logical_map, list) + if _LOG_INPUT: + logger.info( + "update_expert_weights_single_layer " + f"{[x.shape for x in routed_experts_weights]=} " + f"{[x.shape for x in temp_buffers]=} " + f"{old_physical_to_logical_map=} " + f"{new_physical_to_logical_map=} " + f"{num_local_physical_experts=} " + f"{num_gpu_per_node=} " + f"{rank=} " + f"{world_size=} " + ) + output_logs = [] if debug else None num_physical_experts = len(old_physical_to_logical_map)