From 3199b3091db9bf51f348f8ec0a08005fd50dd473 Mon Sep 17 00:00:00 2001
From: Jiacheng Huang <huangjiacheng0709@outlook.com>
Date: Fri, 7 Aug 2026 00:57:53 +0800
Subject: [PATCH] refactor(ops): migrate topksoftmax to canonical InfiniOps API

---
 .../ops/topksoftmax/topksoftmax_infiniops.cc  | 29 ++++++++++++++-----
 submodules/InfiniOps                          |  2 +-
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc b/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc
index 4a4679417..9cb7553ba 100644
--- a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc
+++ b/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc
@@ -3,15 +3,16 @@
 #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
@@ -19,7 +20,17 @@ struct PlannedMeta {
 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) {
@@ -27,14 +38,16 @@ void run(void *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) {
diff --git a/submodules/InfiniOps b/submodules/InfiniOps
index e733e325b..1a712d0c2 160000
--- a/submodules/InfiniOps
+++ b/submodules/InfiniOps
@@ -1 +1 @@
-Subproject commit e733e325b4af776261419e4dc7e1f9386c05fa71
+Subproject commit 1a712d0c25a0cc0a369254e7e0a5056824a82552