Merge pull request #109 from raindaywhu/dev_whq_eplb1

add function take_update_info_from_eplb_process

Author: wanghanqingLYT

vllm_ascend/eplb/eplb_updator.py

@@ -118,10 +118,6 @@ def wakeup_eplb_worker(self):
         self.planner_block_queue.put(1)
 
     def forward_before(self):
-        # Batch after eplb process being triggered, get update info provided by eplb process
-        if self.get_update_info_flag():
-            self.update_info_all = self.block_update_queue.get()
-
         if self.update_expert_weight_flag():
             (expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = self.update_info_all.pop(0)
             rank_id = torch.distributed.get_rank()
@@ -137,6 +133,12 @@ def forward_before(self):
             self.reqs = []
             self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
 
+    def take_update_info_from_eplb_process(self):
+        # Batch after eplb process being triggered, get update info provided by eplb process
+        if self.get_update_info_flag():
+            self.update_info_all = self.block_update_queue.get()
+
+
     def forward_end(self):
         if self.wakeup_eplb_worker_flag():
             moe_load = self.compute_and_set_moe_load(is_clear=True)

vllm_ascend/worker/model_runner_v1.py

@@ -1246,6 +1246,9 @@ def execute_model(
              sample_indices) = (self._process_reqs(scheduler_output,
                                                    intermediate_tensors))
 
+            if self.dynamic_eplb:
+                self.eplb_updator.take_update_info_from_eplb_process()
+
         with ProfileExecuteDuration().capture_async("post process"):
             logits = self.model.compute_logits(hidden_states[sample_indices],
                                                None)
@@ -1547,6 +1550,7 @@ def _dummy_run(
                 if is_profile_run and self.dynamic_eplb:
                     self.model.clear_all_moe_loads()
                 if not is_compile and not is_profile_run and self.dynamic_eplb:
+                    self.eplb_updator.take_update_info_from_eplb_process()
                     self.eplb_updator.forward_end()
                 return hidden_states