Merge pull request #101 from raindaywhu/br_main_into_eplb_wjh

optimize calculate moe load

Author: qmkakaxi

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -38,7 +38,7 @@ def __init__(self, model, **args):
         self.num_moe_layers = self.model.config.num_hidden_layers - self.num_dense_layers
         self.global_expert_num = self.model.config.n_routed_experts
 
-        
+
         # TODO: init self.expert_weight_names depending on different model types, only deepseek v3 w8a8 is supported here
         self.expert_weight_names = ["w13_weight", "w2_weight", "w13_weight_scale", "w13_weight_offset",
             "w2_weight_scale", "w2_weight_offset"]
@@ -62,6 +62,8 @@ def __init__(self, model, **args):
             self.log2phy_map_per_layer[self.num_dense_layers + layer_idx] =\
                 self.model.get_log2phy_map(self.num_dense_layers + layer_idx)
 
+        self.all_topk_ids = []
+
     def init_buffer_tensor(self, num_buffer_tensor):
         for name in self.expert_weight_names:
             complete_name = "model.layers." + str(self.num_dense_layers) + ".mlp.experts." + name
@@ -82,39 +84,36 @@ def init_expert_param_per_layer(self):
                         for name in self.expert_weight_names]
                 )
 
-    def get_rank_expert_workload(
-        self,
-        num_moe_layers: int,
-        dummy_run = False
-    ) -> torch.Tensor:
-
-        all_topk_ids = [self.model.get_topk_ids(i) for i in range(num_moe_layers)]
-        stacked = torch.stack(all_topk_ids, dim=0)      
-        L, B, K = stacked.shape
-        N = B * K
-        device = stacked.device
-        G = self.global_expert_num
+    def collect_topk_ids(self, dummy_run=False):
+        if dummy_run:
+            return 
+        self.all_topk_ids.append(self.model.get_all_topk_ids(self.num_moe_layers))
 
-        if not hasattr(self, "cum_moe_load") or self.cum_moe_load is None:
-            self.cum_moe_load = torch.zeros((L, G),
-                                            dtype=torch.int64,
-                                            device=device)
+    def get_rank_expert_workload(self) -> torch.Tensor:
 
-        if dummy_run:
-            return self.cum_moe_load
+        device = self.all_topk_ids[0][0].device
+        flat_list_per_layer = [[] for _ in range(self.num_moe_layers)]  
 
-        ids1d = stacked.view(-1).to(torch.int64)       
+        for period_data in self.all_topk_ids:     
+            for l in range(self.num_moe_layers):  
+                t = period_data[l]       
+                flat_list_per_layer[l].append(t.reshape(-1))  
 
-        row_idx = torch.arange(L, device=device).repeat_interleave(N) 
+        index_2d = torch.nn.utils.rnn.pad_sequence(
+            [torch.cat(flat_list_per_layer[l]) for l in range(self.num_moe_layers)],
+            batch_first=True, padding_value=-1     
+        ).to(device)               
 
-        combined = row_idx * G + ids1d                     
+        mask = index_2d != -1
+        index_2d = index_2d.masked_select(mask).reshape(self.num_moe_layers, -1)
+        src_2d   = torch.ones_like(index_2d, dtype=torch.int64)
 
-        counts = torch.bincount(combined, minlength=L * G)  
-        workload = counts.view(L, G)                       
+        moe_load = torch.zeros((self.num_moe_layers),  self.global_expert_num,
+                            dtype=torch.int64, device=device)
+        moe_load.scatter_add_(dim=1, index=index_2d, src=src_2d)
 
-        self.cum_moe_load.add_(workload)
-    
-        return self.cum_moe_load
+        self.all_topk_ids = []
+        return moe_load
 
     def get_init_expert_map(self, num_moe_layers):
         expert_map = self.model.get_all_expert_map(num_moe_layers)

vllm_ascend/eplb/eplb_updator.py

@@ -57,7 +57,7 @@ def init_eplb(self, expert_map_path):
         self.cur_iterations: torch.int64 = 0
 
         self.wait_worker_iterations: torch.int64 = 0
-        self.num_wait_worker_iterations: torch.int64 = 10
+        self.num_wait_worker_iterations: torch.int64 = 20
 
         self.planner_block_queue = Queue()
         self.block_update_queue = Queue(maxsize=1)
@@ -131,11 +131,11 @@ def forward_before(self):
         self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
 
     def forward_end(self,dummy_run=False):
-        self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
+        self.adaptor.collect_topk_ids(dummy_run)
         if not self.update_in_flight:
             load_gather_iteration, update_iteration = self.get_update_iteration()
             if load_gather_iteration:
-                moe_load = self.compute_and_set_moe_load(dummy_run)
+                moe_load = self.compute_and_set_moe_load()
             if update_iteration:
                 self.wakeup_eplb_worker()
                 self.update_in_flight = True
@@ -148,7 +148,7 @@ def forward_end(self,dummy_run=False):
         self.eplb_loader.update_expert_map_and_weight(self.reqs, self.redundant_enable)
 
     def compute_and_set_moe_load(self,dummy_run=False):
-        local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
+        local_load = self.adaptor.get_rank_expert_workload()
         self._gather_buffer = None
         if dist.is_initialized():
             self.world_size = dist.get_world_size()
@@ -173,7 +173,7 @@ def compute_and_set_moe_load(self,dummy_run=False):
     def warm_up_eplb(self):
 
         self.get_init_expert_map()
-        
+        self.adaptor.collect_topk_ids(dummy_run=False)
         self.compute_and_set_moe_load()
 
         src_tensor = torch.empty((1,), device=self.device)

vllm_ascend/models/deepseek_v2.py

@@ -772,5 +772,12 @@ def get_all_expert_map(self,num_moe_layers):
     def get_topk_ids(self,layer_id):
         return self.model.layers[layer_id+3].mlp.experts.topk_ids
 
+    def get_all_topk_ids(self,num_moe_layers):
+        all_topk_id = []
+        for layer_id in range(num_moe_layers):
+            load_tensor = self.get_topk_ids(layer_id)  
+            all_topk_id.append(load_tensor)
+        return all_topk_id
+
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass