refactor

Author: yangcheng (AJ)

vllm_ascend/ops/fused_moe.py

@@ -45,6 +45,8 @@
     data_parallel_reduce_scatter
 from vllm_ascend.distributed.parallel_state import get_mc2_group
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
+from vllm_ascend.ops.moe_layer.select_experts import UnquantizedSelectExperts
+from vllm_ascend.ops.moe_layer.config import  SelectExpertConfig
 from vllm_ascend.torchair.utils import npu_stream_switch, npu_wait_tensor
 from vllm_ascend.utils import (AscendSocVersion, dispose_tensor,
                                get_all_reduce_merge_state,
@@ -1034,6 +1036,8 @@ def __init__(self, moe: FusedMoEConfig = None):
         except AttributeError:
             self.moe_all_to_all_group_name = None
 
+        self.select_experts = UnquantizedSelectExperts()
+
     def process_weights_after_loading(self, layer):
         super(UnquantizedFusedMoEMethod,
               self).process_weights_after_loading(layer)
@@ -1065,41 +1069,7 @@ def apply(
         **kwargs,
     ) -> torch.Tensor:
 
-        is_deepseek_v3_r1 = global_num_experts == 256
-        # NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern
-        if is_deepseek_v3_r1:
-            topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
-                router_logits,
-                k=top_k,  # topk当前写8
-                bias=e_score_correction_bias,
-                k_group=topk_group,  # fix: 4
-                group_count=num_expert_group,  # fix 8
-                group_select_mode=1,  # 0: group中的最大; 1: topk2.sum(fix)
-                renorm=0,  # 0: softmax->topk(fix); 1: topk->softmax
-                norm_type=1,  # 0: softmax; 1: sigmoid(fix)
-                # out_flag=False, # todo new api; 第三个输出是否输出
-                # y2_flag=False, # old api; 第三个输出是否输出
-                routed_scaling_factor=1,
-                eps=float(1e-20))
-        elif SELECT_GATING_TOPK_SOTFMAX_EXPERTS:
-            topk_weights, topk_ids = select_gating_top_k_softmax_experts(
-                hidden_states=x,
-                router_logits=router_logits,
-                top_k=top_k,
-                renormalize=renormalize)
-        else:
-            topk_weights, topk_ids = select_experts(
-                hidden_states=x,
-                router_logits=router_logits,
-                top_k=top_k,
-                use_grouped_topk=use_grouped_topk,
-                renormalize=renormalize,
-                topk_group=topk_group,
-                num_expert_group=num_expert_group,
-                custom_routing_function=custom_routing_function,
-                scoring_func=scoring_func,
-                e_score_correction_bias=e_score_correction_bias,
-            )
+        topk_weights, topk_ids = self.select_experts(router_logits, x)
 
         topk_weights = topk_weights.to(x.dtype)
         # this is a naive implementation for experts load balance so as
@@ -1268,6 +1238,20 @@ def __init__(
             in_dtype=params_dtype,
             quant_config=quant_config)
 
+        select_experts_dict = {
+            'top_k' : top_k,
+            'e_score_correction_bias' : e_score_correction_bias,
+            'topk_group' : topk_group,
+            'num_expert_group' : num_expert_group,
+            'custom_routing_function' : custom_routing_function,
+            'scoring_func' : scoring_func,
+            'global_num_experts' : self.global_num_experts,
+            'use_grouped_topk' : use_grouped_topk,
+            'renormalize' : renormalize,
+        }
+
+        SelectExpertConfig(select_experts_dict)
+
         if quant_config is None:
             self.quant_method = AscendUnquantizedFusedMoEMethod(moe)
         else:

vllm_ascend/ops/moe_layer/config.py

@@ -0,0 +1,8 @@
+
+class SelectExpertConfig:
+    def __init__(self, config):
+        self.config = config
+
+    @staticmethod
+    def get_config():
+        return self.config

vllm_ascend/ops/moe_layer/select_experts.py

@@ -0,0 +1,79 @@
+from abc import ABC, abstractmethod
+import torch_npu
+from vllm_ascend.ops.fused_moe import select_experts
+import vllm_ascend.envs as envs_ascend
+SELECT_GATING_TOPK_SOTFMAX_EXPERTS: bool = envs_ascend.SELECT_GATING_TOPK_SOTFMAX_EXPERTS
+
+
+class BaseSelectExperts(ABC):
+
+    def __init__(self):
+        need_param = SelectExpertConfig.get_config
+        self.top_k = need_param["top_k"]
+        self.e_score_correction_bias = need_param["e_score_correction_bias"]
+        self.topk_group = need_param["topk_group"]
+        self.num_expert_group = need_param["num_expert_group"]
+        self.custom_routing_function = need_param["custom_routing_function"]
+        self.scoring_func = need_param["scoring_func"]
+        self.global_num_experts = need_param["global_num_experts"]
+        self.use_grouped_topk = need_param['use_grouped_topk']
+        self.renormalize = need_param['renormalize']
+
+    def forward(self, router_logits: torch.Tensor, x: torch.Tensor):
+        if self.global_num_experts == 256:
+            topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
+                router_logits,
+                k=self.top_k,  # topk当前写8
+                bias=self.e_score_correction_bias,
+                k_group=self.topk_group,  # fix: 4
+                group_count=self.num_expert_group,  # fix 8
+                group_select_mode=1,  # 0: group中的最大; 1: topk2.sum(fix)
+                renorm=0,  # 0: softmax->topk(fix); 1: topk->softmax
+                norm_type=1,  # 0: softmax; 1: sigmoid(fix)
+                # out_flag=False, # todo new api; 第三个输出是否输出
+                # y2_flag=False, # old api; 第三个输出是否输出
+                routed_scaling_factor=1,
+                eps=float(1e-20))
+        else:
+            topk_weights, topk_ids = select_experts(
+                hidden_states=x,
+                router_logits=router_logits,
+                top_k=self.top_k,
+                use_grouped_topk=self.use_grouped_topk,
+                renormalize=self.renormalize,
+                topk_group=self.topk_group,
+                num_expert_group=self.num_expert_group,
+                custom_routing_function=self.custom_routing_function,
+                scoring_func=self.scoring_func,
+                e_score_correction_bias=self.e_score_correction_bias,
+            )
+        return topk_weights, topk_ids
+
+
+class UnquantizedSelectExperts(BaseSelectExperts):
+    def __init__(self):
+        super().__init__()
+        
+    def forward(self, router_logits: torch.Tensor, x: torch.Tensor):
+        if SELECT_GATING_TOPK_SOTFMAX_EXPERTS:
+            topk_weights, topk_ids = select_gating_top_k_softmax_experts(
+                hidden_states=x,
+                router_logits=router_logits,
+                top_k=self.top_k,
+                renormalize=self.renormalize)
+        else:
+            topk_weights, topk_ids = super().forward(router_logits, x)
+
+        return topk_weights, topk_ids
+
+
+class QuantizedSelectExperts(BaseSelectExperts):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, router_logits: torch.Tensor, x: torch.Tensor):
+
+        return  super().forward(router_logits, x)
+
+
+

vllm_ascend/quantization/w8a8_dynamic.py

@@ -27,6 +27,7 @@
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
 from vllm_ascend.distributed.parallel_state import get_mc2_group
+from vllm_ascend.ops.moe_layer.select_experts import QuantizedSelectExperts
 from vllm_ascend.ops.fused_moe import select_experts
 from vllm_ascend.torchair.utils import npu_stream_switch, npu_wait_tensor
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, AscendSocVersion,
@@ -766,6 +767,8 @@ def __init__(self):
         except AttributeError:
             self.moe_all_to_all_group_name = ""
 
+        self.select_experts = QuantizedSelectExperts()
+
     @staticmethod
     def get_weight(num_experts: int, intermediate_size_per_partition: int,
                    hidden_sizes: int,
@@ -835,36 +838,8 @@ def apply(
         assert router_logits.shape[
             1] == global_num_experts, "Number of global experts mismatch"
 
-        is_deepseek_v3_r1 = global_num_experts == 256
-
-        # NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern
-        if is_deepseek_v3_r1:
-            topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
-                router_logits,
-                k=top_k,  # topk当前写8
-                bias=e_score_correction_bias,
-                k_group=topk_group,  # fix: 4
-                group_count=num_expert_group,  # fix 8
-                group_select_mode=1,  # 0: group中的最大; 1: topk2.sum(fix)
-                renorm=0,  # 0: softmax->topk(fix); 1: topk->softmax
-                norm_type=1,  # 0: softmax; 1: sigmoid(fix)
-                # out_flag=False, # todo new api; 第三个输出是否输出
-                # y2_flag=False, # old api; 第三个输出是否输出
-                routed_scaling_factor=1,
-                eps=float(1e-20))
-        else:
-            topk_weights, topk_ids = select_experts(
-                hidden_states=x,
-                router_logits=router_logits,
-                top_k=top_k,
-                use_grouped_topk=use_grouped_topk,
-                renormalize=renormalize,
-                topk_group=topk_group,
-                num_expert_group=num_expert_group,
-                custom_routing_function=custom_routing_function,
-                scoring_func=scoring_func,
-                e_score_correction_bias=e_score_correction_bias,
-            )
+
+        topk_weights, topk_ids = self.select_experts(router_logits, x)
 
         fused_moe_state = get_forward_context().fused_moe_state
         shared_gate_up, shared_dequant_scale = None, None