fix CI issues

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

Author: Yuxiao-Xu

vllm_ascend/ops/expert_load_balancer.py

@@ -1,15 +1,17 @@
 import json
 import random
-from typing import List, Dict
+from typing import Dict, List
 
 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:
@@ -43,19 +45,22 @@ def generate_index_dicts(tensor_2d):
 
     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)
+            (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])
+            self.expert_map_tensor[layer_id])
         result_dict: Dict[int, List[int]] = {}
         for idx, value in enumerate(concatenated):
             key = value.item()
@@ -64,7 +69,8 @@ def generate_log2phy_expert_map(self, layer_id):
             result_dict[key].append(idx)
 
         log2phy_map = torch.full((self.ranks_num, self.global_expert_num),
-                                                     -1, dtype=torch.int32)
+                                 -1,
+                                 dtype=torch.int32)
         for rank in range(self.ranks_num):
             for key in result_dict:
                 indices_in_concat = result_dict[key]
@@ -79,7 +85,7 @@ 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())
+            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
 
@@ -88,8 +94,9 @@ 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
 
 
@@ -99,4 +106,3 @@ def get_global_redundant_expert_num(self):
 #     print(rank_placement_map)
 #     rank_phy2log_map = expert_load_balancer.get_rank_log2phy_map(1, 0)
 #     print(rank_phy2log_map)
-

vllm_ascend/ops/fused_moe.py

@@ -959,6 +959,7 @@ def apply(
 class AscendFusedMoE(FusedMoE):
 
     moe_counter = -1
+
     def __init__(
         self,
         num_experts: int,  # Global number of experts
@@ -1025,22 +1026,22 @@ def __init__(
         self.log2phy = None
         self.global_redundant_expert_num = 0
 
-
         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", 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, 
+                                                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.moe_instance_id,
+                get_ep_group().rank_in_group)
             self.global_redundant_expert_num = \
                         expert_load_balancer.get_global_redundant_expert_num()
         else:

vllm_ascend/quantization/w8a8_dynamic.py

@@ -32,8 +32,7 @@
 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,
@@ -465,13 +464,14 @@ def fused_experts(hidden_states: torch.Tensor,
         group_list_type = 0
 
     # `hidden_states` will be disposed in the `apply_mlp` function
-    hidden_states = apply_mlp(hidden_states,
-                              w1,
-                              w1_scale,
-                              w2,
-                              w2_scale,
-                              expert_tokens,
-                              group_list_type=group_list_type)
+    hidden_states = apply_mlp(
+        hidden_states,
+        w1,
+        w1_scale,
+        w2,
+        w2_scale,
+        expert_tokens,
+        group_list_type=group_list_type)
 
     if expert_map is not None:
         hidden_states.mul_(sorted_weights.unsqueeze(1))
@@ -722,19 +722,19 @@ def apply(
             # 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,
-                                              w1=layer.w13_weight,
-                                              w1_scale=layer.w13_weight_scale,
-                                              w2=layer.w2_weight,
-                                              w2_scale=layer.w2_weight_scale,
-                                              topk_weights=topk_weights,
-                                              topk_ids=topk_ids,
-                                              top_k=top_k,
-                                              expert_map=expert_map,
-                                              ep_group=self.ep_group,
-                                              log2phy=log2phy,
-                                              global_redundant_expert_num=global_redundant_expert_num,
-                                              )
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w1_scale=layer.w13_weight_scale,
+                w2=layer.w2_weight,
+                w2_scale=layer.w2_weight_scale,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                top_k=top_k,
+                expert_map=expert_map,
+                ep_group=self.ep_group,
+                log2phy=log2phy,
+                global_redundant_expert_num=global_redundant_expert_num,
+                )
 
     def process_weights_after_loading(self, layer):
         if self.transpose_weight: