fix bug in pack update info

Author: qmkakaxi

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -343,19 +343,23 @@ def pack_update_info(self, update_info_generator):
 
             maps.append(new_expert_map[self.rank_id])
 
-            log2phy_map = ExpertMapUtils.generate_log2phy_map(new_expert_map) if self.redundant_enable else None
-            log2phy_all.append(log2phy_map)
+            if self.redundant_enable is not None:
+                log2phy_map = ExpertMapUtils.generate_log2phy_map(new_expert_map) 
+                log2phy_all.append(log2phy_map)
 
             layer_ids.append(layer_id)
 
         # 把 list of Tensor 堆成一个大 Tensor
-        stacked_maps      = torch.stack(maps,      dim=0)  # [N, ...]
-        stacked_log2phy   = torch.stack(log2phy_all, dim=0)  # [N, ...]
-        layer_id_tensor   = torch.as_tensor(layer_ids, dtype=torch.int64)  # [N]
+        stacked_maps      = torch.stack(maps,      dim=0)  
+        layer_id_tensor   = torch.as_tensor(layer_ids, dtype=torch.int64)  
+        stacked_maps.share_memory_()
+        layer_id_tensor.share_memory_()
 
-        # 跨进程零拷贝
-        for t in (stacked_maps, stacked_log2phy, layer_id_tensor):
-            t.share_memory_()
+        if self.redundant_enable:
+            stacked_log2phy = torch.stack(log2phy_all, dim=0)
+            stacked_log2phy.share_memory_()
+        else:
+            stacked_log2phy = None
 
         return send_all, recv_all, stacked_maps, stacked_log2phy, layer_id_tensor
 
@@ -375,7 +379,7 @@ def __init__(self, shared_dict, planner_q, block_update_q, redundant_enable, pol
         self.redundant_enable = redundant_enable
 
         # Create EplbWorker instance
-        self.worker = EplbWorker(self.shared_dict, self.policy_type, self.enable_d2d)
+        self.worker = EplbWorker(self.shared_dict, self.policy_type, self.enable_d2d, self.redundant_enable)
 
 
     def worker_process(self, planner_q, block_update_q):
@@ -387,17 +391,12 @@ def worker_process(self, planner_q, block_update_q):
 
                 planner_q.get()
 
-                update_info_generator = self.worker.do_update()
-                update_info_list = []
-
-                for (send_info , recv_info , new_expert_map, layer_id) in update_info_generator:
-                    log2phy_map = ExpertMapUtils.generate_log2phy_map(new_expert_map) if self.redundant_enable else None
-                    update_info_list.append((send_info , recv_info , new_expert_map, log2phy_map, layer_id))
+                packed_update_info = self.worker.do_update()
 
                 while True:
                     if not block_update_q.empty():
                         continue
-                    block_update_q.put(update_info_list)
+                    block_update_q.put(packed_update_info)
                     break
 
             except Exception as e:

vllm_ascend/eplb/eplb_updator.py

@@ -75,7 +75,6 @@ def init_eplb(self, redundant_enable):
 
         logger.info(f"[ModelRunner] Launched EPLB process (pid={self.eplb_process.pid})")
 
-
     def get_update_iteration(self):
         self.cur_iterations = self.cur_iterations + 1
         return self.cur_iterations % self.num_iterations == 0
@@ -191,16 +190,25 @@ def unpack_update_batch(self, packed_update_info):
         """
         send_all, recv_all, stacked_maps, stacked_log2phy, layer_id_tensor = packed_update_info
 
-        # 拆分 Tensor，得到 N 个张量的 tuple
-        maps = stacked_maps.unbind(0)
-        log2phy = stacked_log2phy.unbind(0)
-
-        # 把 layer_id_tensor 转成 Python int 列表
+        maps     = stacked_maps.unbind(0)
         layer_ids = layer_id_tensor.tolist()
 
+        if self.redundant_enable:
+            log2phy_list = stacked_log2phy.unbind(0)
+        else:
+            log2phy_list = [None] * len(maps)
+
+        _zip = zip
+        _send = send_all
+        _recv = recv_all
+        _maps = maps
+        _l2p  = log2phy_list
+        _lids = layer_ids
+
         recovered = [
-            (s, r, m, l, lid)
-            for s, r, m, l, lid in zip(send_all, recv_all, maps, log2phy, layer_ids)
+            (_s, _r, _m, _lp, _lid)
+            for _s, _r, _m, _lp, _lid
+            in _zip(_send, _recv, _maps, _l2p, _lids)
         ]
         return recovered