relocate the code from the worker_runner to the server side.

Author: lt

vllm_ascend/eplb/eplb_updator.py

@@ -234,12 +234,12 @@ def unpack_update_batch(self, packed_update_info):
         ]
         return recovered
 
-    def get_expert_load(self) -> torch.Tensor:
+    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
-        load_info, _ = ExpertMapUtils.global2local_load(moe_load, expert_maps, num_local_experts)
-        return load_info
+        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

@@ -90,7 +90,6 @@ def global2local_load(self,
         placement: torch.Tensor,
         E_local: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
-
         L, G, _ = placement.shape
         device = placement.device
 

vllm_ascend/worker/model_runner_v1.py

@@ -1588,7 +1588,7 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
-    def do_get_expert_load(self) ->  torch.Tensor:
+    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,10 +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) -> torch.Tensor:
-        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)