[AMD][Bugfix][EPLB] Fix elastic EP scaling accuracy on ROCm (#47206)
Signed-off-by: okorzh <okorzh@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Oxana Korzh committed
Jul 24, 2026 at 14:35 UTC
2279575cd9caeb9174c7ec11a94a34af40903f4d
2 files changed
+89
-21
vllm/distributed/eplb/eplb_state.py
+74
-16
@@ -46,6 +46,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
46
from vllm.distributed.utils import StatelessProcessGroup
47
from vllm.logger import init_logger
48
from vllm.model_executor.models.interfaces import MixtureOfExperts
49
+from vllm.platforms import current_platform
50
51
from .async_worker import start_async_worker
52
from .eplb_communicator import EplbCommunicator, create_eplb_communicator
@@ -825,23 +826,80 @@ class EplbState:
826
eplb_model_state.physical_to_logical_map.cpu(),
827
)
828
828
- # Update expert weights
829
- rearrange_expert_weights_inplace(
830
- eplb_model_state.physical_to_logical_map,
831
- new_physical_to_logical_map,
832
- eplb_model_state.model.expert_weights,
833
- eplb_model_state.expert_buffer,
834
- ep_group,
835
- eplb_model_state.communicator,
836
- is_profile,
837
- rank_mapping,
838
- )
839
-
840
- if not is_profile:
841
- _commit_eplb_maps(
842
- eplb_model_state,
843
- new_physical_to_logical_map=new_physical_to_logical_map,
829
+ skip_rearrange = False
830
+ if (
831
+ current_platform.is_rocm()
832
+ and not is_profile
833
+ and rank_mapping is None
834
+ and bool((eplb_model_state.physical_to_logical_map >= 0).all())
835
+ ):
836
+ logical_loads = global_expert_load_window.float()
837
+ ep_size = ep_group.size()
838
+
839
+ def rank_load_imbalance(
840
+ mapping: torch.Tensor,
841
+ logical_loads: torch.Tensor = logical_loads,
842
+ ep_size: int = ep_size,
843
+ ) -> float:
844
+ mapping = mapping.to(
845
+ device=logical_loads.device,
846
+ dtype=torch.long,
847
+ )
848
+ replica_counts = torch.zeros_like(logical_loads)
849
+ replica_counts.scatter_add_(
850
+ dim=1,
851
+ index=mapping,
852
+ src=torch.ones_like(mapping, dtype=logical_loads.dtype),
853
+ )
854
+ loads_per_replica = torch.gather(
855
+ logical_loads / replica_counts.clamp_min(1),
856
+ dim=1,
857
+ index=mapping,
858
+ )
859
+ loads_per_rank = loads_per_replica.reshape(
860
+ logical_loads.shape[0], ep_size, -1
861
+ ).sum(dim=(0, 2))
862
+ mean_load = loads_per_rank.mean()
863
+ if mean_load == 0:
864
+ return 1.0
865
+ return (loads_per_rank.max() / mean_load).item()
866
+
867
+ current_imbalance = rank_load_imbalance(
868
+ eplb_model_state.physical_to_logical_map
869
)
870
+ proposed_imbalance = rank_load_imbalance(
871
+ new_physical_to_logical_map
872
+ )
873
+ relative_improvement = (
874
+ current_imbalance - proposed_imbalance
875
+ ) / current_imbalance
876
+ skip_rearrange = relative_improvement < 0.05
877
+ if skip_rearrange and is_main_rank:
878
+ logger.info(
879
+ "[EPLB] Skip rearrange: imbalance %.4f -> "
880
+ "%.4f (no material gain)",
881
+ current_imbalance,
882
+ proposed_imbalance,
883
+ )
884
+
885
+ if not skip_rearrange:
886
+ # Update expert weights
887
+ rearrange_expert_weights_inplace(
888
+ eplb_model_state.physical_to_logical_map,
889
+ new_physical_to_logical_map,
890
+ eplb_model_state.model.expert_weights,
891
+ eplb_model_state.expert_buffer,
892
+ ep_group,
893
+ eplb_model_state.communicator,
894
+ is_profile,
895
+ rank_mapping,
896
+ )
897
+
898
+ if not is_profile:
899
+ _commit_eplb_maps(
900
+ eplb_model_state,
901
+ new_physical_to_logical_map=new_physical_to_logical_map,
902
+ )
903
904
if is_main_rank:
905
assert start_event is not None
vllm/model_executor/layers/fused_moe/routed_experts.py
+15
-5
@@ -233,17 +233,27 @@ class RoutedExperts(PluggableLayer):
233
# Update local attributes from ExpertMapManager
234
self.local_num_experts = self.expert_map_manager.local_num_experts
235
self.expert_placement_strategy = self.expert_map_manager.placement_strategy
236
- self.register_buffer("_expert_map", self.expert_map_manager.expert_map)
237
- self.register_buffer("expert_mask", self.expert_map_manager.expert_mask)
236
+ self.register_buffer(
237
+ "_expert_map", self.expert_map_manager.expert_map, persistent=False
238
+ )
239
+ self.register_buffer(
240
+ "expert_mask", self.expert_map_manager.expert_mask, persistent=False
241
+ )
242
243
# Get routing tables from ExpertMapManager
244
routing_tables = self.expert_map_manager.routing_tables
245
if routing_tables is not None:
246
# Register routing tables as buffers for this layer
247
global_to_physical, physical_to_global, local_global = routing_tables
244
- self.register_buffer("expert_global_to_physical", global_to_physical)
245
- self.register_buffer("expert_physical_to_global", physical_to_global)
246
- self.register_buffer("expert_local_to_global", local_global)
248
+ self.register_buffer(
249
+ "expert_global_to_physical", global_to_physical, persistent=False
250
+ )
251
+ self.register_buffer(
252
+ "expert_physical_to_global", physical_to_global, persistent=False
253
+ )
254
+ self.register_buffer(
255
+ "expert_local_to_global", local_global, persistent=False
256
+ )
257
258
def _expert_routing_tables(
259
self,