refactor(ops): migrate topksoftmax to canonical InfiniOps API

src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc

@@ -3,38 +3,51 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/topksoftmax_infinilm.h"
+#include "base/topk_softmax.h"
+
+#include <optional>
 
 namespace infinicore::op::topksoftmax_impl::infiniops {
 namespace {
 using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
 struct PlannedMeta {
-    TensorMeta values, indices, x;
-    graph::GraphTensor values_tensor, indices_tensor, x_tensor;
-    size_t topk;
+    TensorMeta values, indices, token_expert_indices, x;
+    graph::GraphTensor values_tensor, indices_tensor, token_expert_indices_tensor, x_tensor;
     int norm;
 };
 } // namespace
 
 void *plan(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) {
     INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(values->device().getType()));
     INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x);
-    return new PlannedMeta{TensorMeta(values), TensorMeta(indices), TensorMeta(x), graph::GraphTensor(values), graph::GraphTensor(indices), graph::GraphTensor(x), topk, norm};
+    auto token_expert_indices = Tensor::empty({x->size(0), topk}, DataType::I32, indices->device());
+    return new PlannedMeta{
+        TensorMeta(values),
+        TensorMeta(indices),
+        TensorMeta(token_expert_indices),
+        TensorMeta(x),
+        graph::GraphTensor(values),
+        graph::GraphTensor(indices),
+        graph::GraphTensor(token_expert_indices),
+        graph::GraphTensor(x),
+        norm};
 }
 
 void run(void *planned_meta) {
     auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
     infini::ops::Config config;
-    infini::ops::TopksoftmaxInfinilm::Call(
+    infini::ops::TopkSoftmax::Call(
         handle,
         config,
         planned->x.tensor(planned->x_tensor),
-        static_cast<int64_t>(planned->topk),
+        std::optional<infini::ops::Tensor>{},
+        std::optional<infini::ops::Tensor>{},
         planned->norm != 0,
         planned->values.tensor(planned->values_tensor),
-        planned->indices.tensor(planned->indices_tensor));
+        planned->indices.tensor(planned->indices_tensor),
+        planned->token_expert_indices.tensor(planned->token_expert_indices_tensor));
 }
 
 void cleanup(void **planned_meta_ptr) {