Revert "[cherry-pick][refactor]support gatingtopk operator generalization (#4050)" (#4352)

This reverts commit c87a77e8b4.

it breaks ops e2e test

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

vllm_ascend/ascend_forward_context.py

@@ -96,7 +96,6 @@ def set_ascend_forward_context(
         ep_size = (get_ep_group().world_size if
                    vllm_config.parallel_config.enable_expert_parallel else 1)
 
-        # fused_moe_state is used in torchair, it will be deleted along with torchair
         is_deepseek_v3_r1 = hasattr(
             vllm_config.model_config.hf_config, 'n_routed_experts'
         ) and vllm_config.model_config.hf_config.n_routed_experts == 256

vllm_ascend/ops/moe/experts_selector.py

@@ -20,6 +20,8 @@ import torch
 import torch_npu
 from vllm.forward_context import get_forward_context
 
+from vllm_ascend.ascend_config import get_ascend_config
+
 
 def select_experts(hidden_states: torch.Tensor,
                    router_logits: torch.Tensor,
@@ -60,20 +62,21 @@ def select_experts(hidden_states: torch.Tensor,
     if weight_prefetch_method:
         weight_prefetch_method.maybe_prefetch_moe_weight_preprocess(
             hidden_states, "gate_up")
-    if custom_routing_function is None:
-        topk_weights, topk_ids = _select_experts_with_fusion_ops(
-            hidden_states=hidden_states,
-            router_logits=router_logits,
-            top_k=top_k,
-            use_grouped_topk=use_grouped_topk,
-            topk_group=topk_group,
-            renormalize=renormalize,
-            e_score_correction_bias=e_score_correction_bias,
-            num_expert_group=num_expert_group,
-            scoring_func=scoring_func,
-            routed_scaling_factor=routed_scaling_factor,
-            global_num_experts=global_num_experts)
-    else:
+    topk_weights, topk_ids = _select_experts_with_fusion_ops(
+        hidden_states=hidden_states,
+        router_logits=router_logits,
+        top_k=top_k,
+        use_grouped_topk=use_grouped_topk,
+        topk_group=topk_group,
+        renormalize=renormalize,
+        e_score_correction_bias=e_score_correction_bias,
+        num_expert_group=num_expert_group,
+        custom_routing_function=custom_routing_function,
+        scoring_func=scoring_func,
+        routed_scaling_factor=routed_scaling_factor,
+        global_num_experts=global_num_experts)
+
+    if topk_weights is None:
         topk_weights, topk_ids = _native_select_experts(
             hidden_states=hidden_states,
             router_logits=router_logits,
@@ -168,34 +171,34 @@ def _select_experts_with_fusion_ops(
         e_score_correction_bias: Optional[torch.Tensor],
         topk_group: Optional[int],
         num_expert_group: Optional[int],
+        custom_routing_function: Optional[Callable] = None,
         scoring_func: str = "softmax",
         routed_scaling_factor=1.0,
         global_num_experts: int = -1):
 
-    if scoring_func == "softmax":
-        norm_type = 0
-        topk_group = 1
-        num_expert_group = 1
-    else:
-        norm_type = 1
-    if e_score_correction_bias is not None and \
-        e_score_correction_bias.dtype != router_logits.dtype:
-        e_score_correction_bias = e_score_correction_bias.to(
-            router_logits.dtype)
-    topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
-        router_logits,
-        k=top_k,
-        bias=e_score_correction_bias,
-        k_group=topk_group,
-        group_count=num_expert_group,
-        group_select_mode=1,  # 0: the maximum in the group; 1: topk2.sum(fix)
-        renorm=0,  # 0: softmax->topk(fix); 1: topk->softmax
-        norm_type=norm_type,  # 0: softmax; 1: sigmoid
-        # out_flag=False, # todo new api; should the third output be output
-        # y2_flag=False, # old api; should the third output be output
-        routed_scaling_factor=1,
-        eps=float(1e-20))
-    if scoring_func == "softmax":
+    topk_weights, topk_ids = None, None
+    # NOTE: now npu_moe_gating_top_k can only support 'group_count=256' pattern
+    global_redundant_expert_num = get_ascend_config().init_redundancy_expert
+    is_deepseek_v3_r1 = global_num_experts - global_redundant_expert_num == 256
+    if is_deepseek_v3_r1:
+        topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
+            router_logits,
+            k=top_k,  # topk currently 8
+            bias=e_score_correction_bias,
+            k_group=topk_group,  # fix: 4
+            group_count=num_expert_group,  # fix 8
+            group_select_mode=
+            1,  # 0: the maximum in the 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; should the third output be output
+            # y2_flag=False, # old api; should the third output be output
+            routed_scaling_factor=1,
+            eps=float(1e-20))
+    if not use_grouped_topk and custom_routing_function is None and scoring_func == "softmax":
+        topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k_softmax(
+            x=router_logits, finished=None, k=top_k)
+        topk_ids = topk_ids.to(torch.int32)
         topk_weights = _renormalize_topk_weights(topk_weights, renormalize)
 
     return topk_weights, topk_ids