[Bug] Fix moe_sum signature (#18440)

Signed-off-by: Bill Nell <bnell@redhat.com>
2025-05-21 01:37:08 -04:00
parent 0c15c2e486
commit 92247c522e
2 changed files with 19 additions and 1 deletions
--- a/csrc/moe/torch_bindings.cpp
+++ b/csrc/moe/torch_bindings.cpp
@@ -10,7 +10,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {

  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
-  m.def("moe_sum(Tensor! input, Tensor output) -> ()");
+  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
  m.impl("moe_sum", torch::kCUDA, &moe_sum);

  // Aligning the number of tokens to be processed by each expert such