fix CI problem

Signed-off-by: Yuxiao-Xu <664988918@qq.com>

Author: Yuxiao-Xu

vllm_ascend/ops/expert_load_balancer.py

@@ -1,12 +1,15 @@
-import torch
 import json
 import random
+from typing import List, Dict
+
+import torch
 
 class ExpertLoadBalancer(object):
     def __init__(self, expert_map_path, global_expert_num):
         self.expert_map_path = expert_map_path
         self.global_expert_num = global_expert_num
-        self.expert_map_tensor, self.layers_num, self.ranks_num = self.expert_file_to_tensor()
+        self.expert_map_tensor, self.layers_num, self.ranks_num = \
+                                           self.expert_file_to_tensor()
 
     def expert_file_to_tensor(self):
         with open(self.expert_map_path, "r") as f:
@@ -39,25 +42,29 @@ def generate_index_dicts(tensor_2d):
         return dict_list
 
     def generate_expert_placement_map(self):
-        expert_placement_map = torch.full((self.layers_num, self.ranks_num, self.global_expert_num),
-                                          -1, dtype=torch.int32)
+        expert_placement_map = torch.full(
+                    (self.layers_num, self.ranks_num, self.global_expert_num),
+                     -1, dtype=torch.int32)
         for layer_id in range(self.layers_num):
             for gpu_id in range(self.ranks_num):
                 e_ids = self.expert_map_tensor[layer_id, gpu_id]
-                expert_placement_map[layer_id, gpu_id, e_ids] = torch.arange(len(e_ids), dtype=torch.int32)
+                expert_placement_map[layer_id, gpu_id, e_ids] = \
+                                  torch.arange(len(e_ids), dtype=torch.int32)
         return expert_placement_map
 
     def generate_log2phy_expert_map(self, layer_id):
         concatenated = torch.flatten(self.expert_map_tensor[layer_id])
-        rank_expert_to_global = self.generate_index_dicts(self.expert_map_tensor[layer_id])
-        result_dict = {}
+        rank_expert_to_global = self.generate_index_dicts(
+                                            self.expert_map_tensor[layer_id])
+        result_dict: Dict[int, List[int]] = {}
         for idx, value in enumerate(concatenated):
             key = value.item()
             if key not in result_dict:
                 result_dict[key] = []
             result_dict[key].append(idx)
 
-        log2phy_map = torch.full((self.ranks_num, self.global_expert_num), -1, dtype=torch.int32)
+        log2phy_map = torch.full((self.ranks_num, self.global_expert_num),
+                                                     -1, dtype=torch.int32)
         for rank in range(self.ranks_num):
             for key in result_dict:
                 indices_in_concat = result_dict[key]
@@ -71,7 +78,8 @@ def generate_log2phy_expert_map(self, layer_id):
     def get_rank_placement_map(self, layer_id, rank_id):
         expert_placement_map = self.generate_expert_placement_map()
         layer_expert_map = expert_placement_map[layer_id]
-        rank_expert_map = layer_expert_map[rank_id].to(torch.npu.current_device())
+        rank_expert_map = layer_expert_map[rank_id].to(
+                                        torch.npu.current_device())
         rank_local_expert_num = torch.sum(torch.ne(rank_expert_map, -1)).item()
         return rank_local_expert_num, rank_expert_map
 
@@ -80,7 +88,8 @@ def get_rank_log2phy_map(self, layer_id, rank_id):
         return layer_log2phy_map[rank_id]
 
     def get_global_redundant_expert_num(self):
-        global_redundant_expert_num = len(self.expert_map_tensor[0][0]) * self.ranks_num - self.global_expert_num
+        global_redundant_expert_num = len(self.expert_map_tensor[0][0]) \
+                                      * self.ranks_num - self.global_expert_num
         return global_redundant_expert_num
 
 

vllm_ascend/ops/fused_moe.py

@@ -1029,15 +1029,20 @@ def __init__(
         vllm_config = get_current_vllm_config()
         expert_map_path = None
         if vllm_config.additional_config:
-            expert_map_path = vllm_config.additional_config.get("expert_map_path", None)
+            expert_map_path = vllm_config.additional_config.get(
+                                            "expert_map_path", None)
         if expert_map_path and os.path.exists(expert_map_path):
             # moe expert load balance
-            expert_load_balancer = ExpertLoadBalancer(expert_map_path, self.global_num_experts)
-            self.local_num_experts, self.expert_map = expert_load_balancer.get_rank_placement_map(
-                self.moe_instance_id, get_ep_group().rank_in_group)
+            expert_load_balancer = ExpertLoadBalancer(expert_map_path,
+                                                      self.global_num_experts)
+            self.local_num_experts, self.expert_map = \
+                                expert_load_balancer.get_rank_placement_map(
+                                                self.moe_instance_id, 
+                                                get_ep_group().rank_in_group)
             self.log2phy = expert_load_balancer.get_rank_log2phy_map(
                 self.moe_instance_id, get_ep_group().rank_in_group)
-            self.global_redundant_expert_num = expert_load_balancer.get_global_redundant_expert_num()
+            self.global_redundant_expert_num = \
+                        expert_load_balancer.get_global_redundant_expert_num()
         else:
             # Create a tensor of size num_experts filled with -1
             self.local_num_experts, self.expert_map = determine_expert_map(

vllm_ascend/quantization/quant_config.py

@@ -331,7 +331,8 @@ def apply(
             layer, x, router_logits, top_k, renormalize, use_grouped_topk,
             global_num_experts, expert_map, topk_group, num_expert_group,
             custom_routing_function, scoring_func, e_score_correction_bias,
-            is_prefill, enable_force_load_balance, log2phy, global_redundant_expert_num, **kwargs)
+            is_prefill, enable_force_load_balance, log2phy,
+            global_redundant_expert_num, **kwargs)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if hasattr(self.quant_method, "process_weights_after_loading"):

vllm_ascend/quantization/w8a8_dynamic.py

@@ -32,7 +32,8 @@
 VLLM_ENABLE_MC2: bool = envs_ascend.VLLM_ENABLE_MC2
 
 
-def apply_mlp(hidden_states: torch.Tensor,
+def apply_mlp(
+              hidden_states: torch.Tensor,
               w1: torch.Tensor,
               w1_scale: torch.Tensor,
               w2: torch.Tensor,
@@ -720,7 +721,8 @@ def apply(
             # according to tp_size before they are feed into fused_moe module.
             # Therefore, all2all is needed no matter how dp/tp is set so as to
             # dispatch/combine tokens.
-            return fused_experts_with_all2all(hidden_states=x,
+            return fused_experts_with_all2all(
+                                              hidden_states=x,
                                               w1=layer.w13_weight,
                                               w1_scale=layer.w13_weight_scale,
                                               w2=layer.w2_weight,