[dlinfer] fix moe op for dlinfer. (#2917)

* fix fused_moe

* refine code.

* refine code.

---------

Co-authored-by: yaofengchen <[email protected]>

Author: Reinerzhou

lmdeploy/pytorch/backends/dlinfer/moe.py

@@ -47,8 +47,8 @@ def forward(self,
                 down_weights: torch.Tensor,
                 expert_list: List[int] = None):
         """forward."""
-        return fused_moe(hidden_states, self.top_k, topk_ids, topk_weights,
-                         gate_up_weights, down_weights)
+        return fused_moe(hidden_states, gate_up_weights, down_weights,
+                         topk_weights, topk_ids, self.top_k, self.renormalize)
 
 
 class DlinferFusedMoEBuilder(FusedMoEBuilder):

lmdeploy/pytorch/kernels/dlinfer/fused_moe.py

@@ -5,12 +5,13 @@
 
 def fused_moe(
     hidden_states: Tensor,
-    top_k: int,
-    topk_ids: Tensor,
-    topk_weights: Tensor,
     gate_up_weights: Tensor,
     down_weights: Tensor,
+    topk_weights: Tensor,
+    topk_ids: Tensor,
+    topk: int,
+    renormalize: bool,
 ):
-    """ascend fused moe."""
-    return ext_ops.fused_moe(hidden_states, top_k, topk_ids, topk_weights,
-                             gate_up_weights, down_weights)
+    """dlinfer fused moe."""
+    return ext_ops.fused_moe(hidden_states, gate_up_weights, down_weights,
+                             topk_weights, topk_ids, topk, renormalize)