From 2bb7605ae197d54fdf04f4364cddfc97962b7fcd Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Sat, 26 Sep 2026 15:05:37 +0300 Subject: [PATCH] [Elastic EP] Fix EPLB load statistics during scaling (#58473) Signed-off-by: Itay Alroy Co-authored-by: Tyler Michael Smith --- vllm/distributed/elastic_ep/elastic_execute.py | 5 +---- vllm/distributed/eplb/eplb_state.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 19fc15b275..2602edf2df 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -550,10 +550,7 @@ class ElasticEPScalingExecutor: is_async_enabled = eplb_state.is_async run_async = async_op and is_async_enabled eplb_state.is_async = run_async - if rank_mapping is None: - eplb_state.rearrange() - else: - eplb_state.rearrange(rank_mapping=rank_mapping) + eplb_state.rearrange(rank_mapping=rank_mapping, use_last_expert_load=True) if not run_async: # Wait for non-blocking expert resharding copies before continuing # the Elastic EP reconfiguration. diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index e4d805b9fe..205226cde4 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -183,6 +183,8 @@ class EplbModelState: See: https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856 """ + last_expert_load: torch.Tensor + """Global logical-expert totals from the last reshuffle; equal demand initially.""" model_name: str model: MixtureOfExperts expert_buffer: list[torch.Tensor] @@ -499,6 +501,11 @@ class EplbState: expert_load_pass=expert_load_pass, expert_load_pass_buffer=expert_load_pass_buffer, expert_load_window=expert_load_window, + last_expert_load=torch.ones( + (model.num_moe_layers, model.num_logical_experts), + dtype=torch.int64, + device=self.device, + ), model_name=model_config.model, model=model, expert_buffer=expert_buffer, @@ -744,6 +751,7 @@ class EplbState: self, is_profile: bool = False, rank_mapping: dict[int, int] | None = None, + use_last_expert_load: bool = False, ) -> torch.Tensor | None: """Rearrange the experts according to the current load. @@ -753,6 +761,7 @@ class EplbState: no memory movement will be performed. Default is False. rank_mapping (dict[int, int] | None): The rank mapping when scaling is done in EEP. + use_last_expert_load: Use the last reshuffle's logical-expert totals. """ ep_group = get_ep_group().device_group @@ -775,6 +784,13 @@ class EplbState: # Map the physical expert load to global logical experts global_expert_load_windows = [] for eplb_model_state in self.model_states.values(): + if use_last_expert_load: + # Distribute rank 0's saved counts through the existing all-reduce. + saved_load = eplb_model_state.last_expert_load + global_expert_load_windows.append( + saved_load.clone() if is_main_rank else torch.zeros_like(saved_load) + ) + continue expert_load_window = eplb_model_state.expert_load_window physical_to_logical = eplb_model_state.physical_to_logical_map invalid_idx = eplb_model_state.model.num_logical_experts @@ -835,6 +851,8 @@ class EplbState: for eplb_model_state, global_expert_load_window in zip( self.model_states.values(), global_expert_load_windows ): + if not is_profile: + eplb_model_state.last_expert_load = global_expert_load_window if not self.is_async or is_profile: # Get new expert mappings for the model. The policy runs on the # host, so the load window and current map have to come back.