Merge branch 'br_main_into_eplb' into dev_whq_eplb2

Author: wanghanqingLYT

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -68,7 +68,7 @@ def do_update(self):
             new_placement = torch.tensor(new_placement)
         self.check_expert_placement(old_placement, new_placement)
         new_expert_maps = self.local2global(new_placement)
-
+        self.update_expert_map(new_expert_maps)
         logger.debug(f"[EPLB Process  new_map differs, performing D2D")
 
         update_info = self.compose_expert_update_info_bipartite(new_expert_maps, self.old_expert_maps)\

vllm_ascend/eplb/eplb_updator.py

@@ -74,8 +74,6 @@ def init_eplb(self, expert_map_path):
             "moe_load": None,
             # 所有的专家表[num_layers, world_size, num_experts]
             "expert_maps": None,
-            # 热度负载信息 [num_layers, world_size, local_num_experts]
-            "load_info": None,
         })
 
         self.eplb = EplbProcess(
@@ -141,6 +139,7 @@ def forward_end(self,dummy_run=False):
             load_gather_iteration, update_iteration = self.get_update_iteration()
             if load_gather_iteration:
                 moe_load = self.compute_and_set_moe_load()
+                self.get_expert_load()
             if update_iteration:
                 self.wakeup_eplb_worker()
                 self.update_in_flight = True
@@ -234,25 +233,11 @@ def unpack_update_batch(self, packed_update_info):
         ]
         return recovered
 
-    def get_expert_load(self) -> str:
-
-        load_info = self.shared_dict["load_info"]  # Tensor [L, W, local_experts_num]
-        L, W, _ = load_info.shape
-
-        expert_load: Dict[str, List[dict]] = {}
-        for c in range(W):
-            layers: List[dict] = []
-            for l in range(L):
-                counts_1d = load_info[l, c]
-
-                layer_val = {
-                    f"expert_{e}": int(v)
-                    for e, v in enumerate(counts_1d.tolist())
-                }
-                layers.append({f"layer_{l}": layer_val})
-            expert_load[f"card_{c}"] = layers
-
-        return {"expert_load": expert_load}
+    def get_expert_load(self) -> tuple:
+        expert_maps = self.shared_dict["expert_maps"]
+        moe_load = self.shared_dict["moe_load"]  # Tensor [L, W, global_experts_num]
+        num_local_experts = expert_maps.max() + 1
+        return  moe_load, expert_maps, num_local_experts
 
     def update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):
         logger.info(f" start update {self.num_expert_load_gather=}, {self.num_iterations}...")

vllm_ascend/eplb/tool/eplb_utils.py

@@ -83,3 +83,32 @@ def global2local(cls,
         pt_local[g_idx, slot_idx] = k_idx
 
         return pt_local
+
+    @classmethod
+    def global2local_load(self,
+        workload: torch.Tensor,
+        placement: torch.Tensor,
+        E_local: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        L, G, _ = placement.shape
+        device = placement.device
+
+        wt_local = torch.full((L, G, E_local),
+                              fill_value=-1,
+                              dtype=workload.dtype,
+                              device=device)
+        pt_local = torch.full((L, G, E_local),
+                              fill_value=-1,
+                              dtype=torch.long,
+                              device=device)
+
+        valid = placement >= 0
+        l_idx, g_idx, k_idx = valid.nonzero(as_tuple=True)
+
+        slot_idx = placement[l_idx, g_idx, k_idx]
+        values = workload[l_idx, g_idx, k_idx]
+
+        wt_local[l_idx, g_idx, slot_idx] = values
+        pt_local[l_idx, g_idx, slot_idx] = k_idx
+
+        return wt_local, pt_local

vllm_ascend/worker/model_runner_v1.py

@@ -1590,7 +1590,7 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
-    def do_get_expert_load(self) -> str:
+    def do_get_expert_load(self) ->  tuple:
         return self.eplb_updator.get_expert_load()
 
     def do_update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):

vllm_ascend/worker/worker_v1.py

@@ -209,11 +209,8 @@ def compile_or_warm_up_model(self) -> None:
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
 
-    def get_expert_load(self) -> str:
-        """ todo 一共几个worker"""
-        moe_load = self.model_runner.do_get_expert_load()
-        return moe_load
-
+    def get_expert_load(self) -> tuple:
+        return self.model_runner.do_get_expert_load()
     def update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):
         self.model_runner.do_update_expert_load_statistical_period(num_expert_load_gather, num_iterations)