[Elastic EP] Fix EPLB load statistics during scaling (#58473)
Signed-off-by: Itay Alroy <ialroy@nvidia.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Itay Alroy committed
Sep 26, 2026 at 15:05 UTC
2bb7605ae197d54fdf04f4364cddfc97962b7fcd
2 files changed
+19
-4
vllm/distributed/elastic_ep/elastic_execute.py
+1
-4
@@ -550,10 +550,7 @@ class ElasticEPScalingExecutor:
550
is_async_enabled = eplb_state.is_async
551
run_async = async_op and is_async_enabled
552
eplb_state.is_async = run_async
553
- if rank_mapping is None:
554
- eplb_state.rearrange()
555
- else:
556
- eplb_state.rearrange(rank_mapping=rank_mapping)
553
+ eplb_state.rearrange(rank_mapping=rank_mapping, use_last_expert_load=True)
554
if not run_async:
555
# Wait for non-blocking expert resharding copies before continuing
556
# the Elastic EP reconfiguration.
vllm/distributed/eplb/eplb_state.py
+18
@@ -183,6 +183,8 @@ class EplbModelState:
183
See:
184
https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856
185
"""
186
+ last_expert_load: torch.Tensor
187
+ """Global logical-expert totals from the last reshuffle; equal demand initially."""
188
model_name: str
189
model: MixtureOfExperts
190
expert_buffer: list[torch.Tensor]
@@ -499,6 +501,11 @@ class EplbState:
501
expert_load_pass=expert_load_pass,
502
expert_load_pass_buffer=expert_load_pass_buffer,
503
expert_load_window=expert_load_window,
504
+ last_expert_load=torch.ones(
505
+ (model.num_moe_layers, model.num_logical_experts),
506
+ dtype=torch.int64,
507
+ device=self.device,
508
+ ),
509
model_name=model_config.model,
510
model=model,
511
expert_buffer=expert_buffer,
@@ -744,6 +751,7 @@ class EplbState:
751
self,
752
is_profile: bool = False,
753
rank_mapping: dict[int, int] | None = None,
754
+ use_last_expert_load: bool = False,
755
) -> torch.Tensor | None:
756
"""Rearrange the experts according to the current load.
757
@@ -753,6 +761,7 @@ class EplbState:
761
no memory movement will be performed. Default is False.
762
rank_mapping (dict[int, int] | None): The rank mapping
763
when scaling is done in EEP.
764
+ use_last_expert_load: Use the last reshuffle's logical-expert totals.
765
766
"""
767
ep_group = get_ep_group().device_group
@@ -775,6 +784,13 @@ class EplbState:
784
# Map the physical expert load to global logical experts
785
global_expert_load_windows = []
786
for eplb_model_state in self.model_states.values():
787
+ if use_last_expert_load:
788
+ # Distribute rank 0's saved counts through the existing all-reduce.
789
+ saved_load = eplb_model_state.last_expert_load
790
+ global_expert_load_windows.append(
791
+ saved_load.clone() if is_main_rank else torch.zeros_like(saved_load)
792
+ )
793
+ continue
794
expert_load_window = eplb_model_state.expert_load_window
795
physical_to_logical = eplb_model_state.physical_to_logical_map
796
invalid_idx = eplb_model_state.model.num_logical_experts
@@ -835,6 +851,8 @@ class EplbState:
851
for eplb_model_state, global_expert_load_window in zip(
852
self.model_states.values(), global_expert_load_windows
853
):
854
+ if not is_profile:
855
+ eplb_model_state.last_expert_load = global_expert_load_window
856
if not self.is_async or is_profile:
857
# Get new expert mappings for the model. The policy runs on the
858
# host, so the load window and current map have to come back.