EPLB moed load collect

Author: qmkakaxi

vllm_ascend/eplb/core/loader/ssd_loader.py

@@ -723,6 +723,39 @@ def forward(
                                    inputs_embeds)
         return hidden_states
 
+    def update_expert_map(self,new_expert_map,layer_id):
+        self.model.layers[layer_id].mlp.experts.update_map(new_expert_map)
+
+    def update_all_expert_map(self, new_expert_map,num_moe_layers):
+        num_moe_layers = len(new_expert_map)
+        for layer_id in range(num_moe_layers):
+            layer_map = new_expert_map[layer_id].to("npu")
+            self.model.layers[3+layer_id].mlp.experts.update_map(layer_map)
+
+    def get_expert_map(self,layer_id):
+        return self.model.layers[layer_id].mlp.experts.get_map()
+
+    def get_all_expert_map(self,num_moe_layers):
+        all_loads = []
+        for layer_id in range(num_moe_layers):
+            load_tensor = self.get_expert_map(3+layer_id)  # (num_experts_per_layer,)
+            all_loads.append(load_tensor)
+
+        return torch.stack(all_loads, dim=0)
+
+    def get_moe_load(self,layer_id):
+        return self.model.layers[layer_id].mlp.experts.get_moe_load()
+
+    def get_all_moe_loads(self, num_moe_layers) -> torch.Tensor:
+        """
+        output: [num_moe_layers, num_experts_per_layer]
+        """
+        all_loads = []
+        for layer_id in range(num_moe_layers):
+            load_tensor = self.get_moe_load(3+layer_id)  # (num_experts_per_layer,)
+            all_loads.append(load_tensor)
+
+        return torch.stack(all_loads, dim=0)
 
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass

vllm_ascend/models/deepseek_v2.py

@@ -991,6 +991,8 @@ def __init__(
 
         AscendFusedMoE.moe_counter += 1
         self.moe_instance_id = AscendFusedMoE.moe_counter
+        self.moe_load = None
+        self.topk_ids =  None
 
         if params_dtype is None:
             params_dtype = torch.get_default_dtype()
@@ -1132,7 +1134,7 @@ def forward(self,
                         hidden_states, router_logits)
 
         # Matrix multiply.
-        hidden_states = self.quant_method.apply(
+        hidden_states, self.topk_ids = self.quant_method.apply(
             layer=self,
             x=hidden_states,
             router_logits=router_logits,
@@ -1152,6 +1154,8 @@ def forward(self,
             global_redundant_expert_num=self.global_redundant_expert_num,
             **kwargs)
 
+        self.calculate_moe_load()
+
         if self.enable_multistream_shared_expert and not is_prefill:
             hidden_states, shared_output = hidden_states
 
@@ -1209,3 +1213,25 @@ def _forward_ms_fused_moe_comp(
             enable_force_load_balance=enable_force_load_balance)
 
         return hidden_states
+        
+    def update_map(self,new_expert_map):
+        self.expert_map = new_expert_map
+
+    def get_map(self):
+        return self.expert_map
+
+    def get_moe_load(self):
+        return self.moe_load
+
+    def calculate_moe_load(self):
+        if self.moe_load is None:
+            self.moe_load = torch.zeros(self.num_experts,
+                                        dtype=torch.int64,
+                                        device=self.topk_ids.device)
+
+        ids     = self.topk_ids.flatten().to(torch.int64)
+
+        ones = torch.ones_like(ids, dtype=torch.int64, device=ids.device)
+        self.moe_load.scatter_add_(0, ids, ones)
+
+        return self.moe_load