fix moe op for dlinfer.

InternLM · Dec 18, 2024 · 4c1be55 · 4c1be55
1 parent 1efed79
commit 4c1be55
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/lmdeploy/pytorch/backends/dlinfer/moe.py b/lmdeploy/pytorch/backends/dlinfer/moe.py
@@ -48,7 +48,7 @@ def forward(self,
                 expert_list: List[int] = None):
         """forward."""
         return fused_moe(hidden_states, self.top_k, topk_ids, topk_weights,
-                         gate_up_weights, down_weights)
+                         gate_up_weights, down_weights, self.renormalize)
 
 
 class DlinferFusedMoEBuilder(FusedMoEBuilder):

diff --git a/lmdeploy/pytorch/kernels/dlinfer/fused_moe.py b/lmdeploy/pytorch/kernels/dlinfer/fused_moe.py
@@ -10,6 +10,7 @@ def fused_moe(
     topk_weights: Tensor,
     gate_up_weights: Tensor,
     down_weights: Tensor,
+    renormalize: bool = False,
 ):
     """ascend fused moe."""
     return ext_ops.fused_moe(hidden_states, top_k, topk_ids, topk_weights,