refactor(ops): migrate InfiniLM adapters to canonical InfiniOps APIs

src/infinicore/ops/conv2d/conv2d_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/conv_infinilm.h"
+#include "base/convolution.h"
 
 #include <optional>
 #include <vector>
@@ -48,15 +48,17 @@ void calculate(Tensor output,
         bias_meta.emplace(bias);
     }
 
-    infini::ops::ConvInfinilm::Call(
+    infini::ops::Convolution::Call(
         handle,
         config,
         input_meta.tensor(input),
         weight_meta.tensor(weight),
         bias_meta ? std::optional<infini::ops::Tensor>{bias_meta->tensor(bias)} : std::nullopt,
-        toInt64Vector(pads, n),
         toInt64Vector(strides, n),
+        toInt64Vector(pads, n),
         toInt64Vector(dilations, n),
+        false,
+        std::vector<int64_t>(n, 0),
         int64_t{1},
         output_meta.tensor(output));
 }

src/infinicore/ops/gelu/gelu_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/gelu_infinilm.h"
+#include "base/gelu.h"
 
 #include <string>
 
@@ -22,7 +22,7 @@ void calculate(Tensor output, Tensor input) {
 
     TensorMeta output_meta(output);
     TensorMeta input_meta(input);
-    infini::ops::GeluInfinilm::Call(
+    infini::ops::Gelu::Call(
         handle,
         config,
         input_meta.tensor(input),

src/infinicore/ops/gelutanh/gelutanh_infiniops.cc

@@ -3,7 +3,9 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/gelutanh_infinilm.h"
+#include "base/gelu.h"
+
+#include <string>
 
 namespace infinicore::op::gelutanh_impl::infiniops {
 namespace {
@@ -20,10 +22,11 @@ void calculate(Tensor output, Tensor input) {
 
     TensorMeta output_meta(output);
     TensorMeta input_meta(input);
-    infini::ops::GelutanhInfinilm::Call(
+    infini::ops::Gelu::Call(
         handle,
         config,
         input_meta.tensor(input),
+        std::string{"tanh"},
         output_meta.tensor(output));
 }
 

src/infinicore/ops/paged_caching/paged_caching_infiniops.cc

@@ -1,36 +1,49 @@
 #include "infinicore/ops/paged_caching.hpp"
 
+#include <string>
+
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/paged_caching_infinilm.h"
+#include "base/reshape_and_cache_flash.h"
 
 namespace infinicore::op::paged_caching_impl::infiniops {
 namespace {
 using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
 struct PlannedMeta {
-    TensorMeta k_cache, v_cache, k, v, slot_mapping;
-    graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, slot_mapping_tensor;
+    TensorMeta k, v, slot_mapping, scale, k_cache, v_cache;
+    graph::GraphTensor k_tensor, v_tensor, slot_mapping_tensor, scale_tensor, k_cache_tensor, v_cache_tensor;
 };
 } // namespace
 
 void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) {
     INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType()));
     INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping);
-    return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping)};
+
+    // The "auto" cache path ignores scales, but the canonical API requires them.
+    auto scale = Tensor::empty({1}, DataType::F32, k_cache->device());
+    auto k_cache_view = k_cache->permute({0, 2, 1, 3});
+    auto v_cache_view = v_cache->permute({0, 2, 1, 3});
+
+    return new PlannedMeta{
+        TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), TensorMeta(scale), TensorMeta(k_cache_view), TensorMeta(v_cache_view),
+        graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping), graph::GraphTensor(scale), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache)};
 }
 
 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::PagedCachingInfinilm::Call(
+    infini::ops::ReshapeAndCacheFlash::Call(
         handle,
         config,
         planned->k.tensor(planned->k_tensor),
         planned->v.tensor(planned->v_tensor),
         planned->slot_mapping.tensor(planned->slot_mapping_tensor),
+        planned->scale.tensor(planned->scale_tensor),
+        planned->scale.tensor(planned->scale_tensor),
+        std::string{"auto"},
         planned->k_cache.tensor(planned->k_cache_tensor),
         planned->v_cache.tensor(planned->v_cache_tensor));
 }

src/infinicore/ops/rearrange/rearrange_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/rearrange_infinilm.h"
+#include "base/copy.h"
 
 namespace infinicore::op::rearrange_impl::infiniops {
 namespace {
@@ -25,17 +25,14 @@ void run(void *planned_meta) {
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
     infini::ops::Config config;
-    infini::ops::RearrangeInfinilm::Call(
-        handle, config, planned->x.tensor(planned->x_tensor), planned->y.tensor(planned->y_tensor));
+    infini::ops::Copy::Call(
+        handle, config, planned->x.tensor(planned->x_tensor), false, planned->y.tensor(planned->y_tensor));
 }
 
 void cleanup(void **planned_meta_ptr) {
     delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
     *planned_meta_ptr = nullptr;
 }
 
-// ReArrange is used by Tensor::copy_from during test result conversion. Keep
-// the InfiniCore/InfiniOp implementation active until this adapter gets the same
-// CUDA shim treatment as Gemm.
 } // namespace infinicore::op::rearrange_impl::infiniops
 #endif

src/infinicore/ops/relu/relu_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/relu_infinilm.h"
+#include "base/relu.h"
 
 namespace infinicore::op::relu_impl::infiniops {
 namespace {
@@ -20,7 +20,7 @@ void calculate(Tensor output, Tensor input) {
 
     TensorMeta output_meta(output);
     TensorMeta input_meta(input);
-    infini::ops::ReluInfinilm::Call(
+    infini::ops::Relu::Call(
         handle,
         config,
         input_meta.tensor(input),

src/infinicore/ops/sigmoid/sigmoid_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/sigmoid_infinilm.h"
+#include "base/sigmoid.h"
 
 namespace infinicore::op::sigmoid_impl::infiniops {
 namespace {
@@ -35,7 +35,7 @@ void run(void *planned_meta) {
     handle.set_stream(context::getStream());
     infini::ops::Config config;
 
-    infini::ops::SigmoidInfinilm::Call(
+    infini::ops::Sigmoid::Call(
         handle,
         config,
         planned->input.tensor(planned->input_tensor),

src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/silu_and_mul_infinilm.h"
+#include "base/silu_and_mul.h"
 
 namespace infinicore::op::silu_and_mul_impl::infiniops {
 namespace {
@@ -25,7 +25,7 @@ void run(void *planned_meta) {
     infini::ops::Handle handle;
     handle.set_stream(context::getStream());
     infini::ops::Config config;
-    infini::ops::SiluAndMulInfinilm::Call(
+    infini::ops::SiluAndMul::Call(
         handle, config, planned->input.tensor(planned->input_tensor), planned->output.tensor(planned->output_tensor));
 }
 

src/infinicore/ops/softmax/softmax_infiniops.cc

@@ -3,7 +3,7 @@
 #ifdef ENABLE_INFINIOPS_API
 #include "../infiniops_impl.hpp"
 
-#include "base/softmax_infinilm.h"
+#include "base/softmax.h"
 
 #include <optional>
 
@@ -22,7 +22,7 @@ void calculate(Tensor output, Tensor input, int axis) {
 
     TensorMeta output_meta(output);
     TensorMeta input_meta(input);
-    infini::ops::SoftmaxInfinilm::Call(
+    infini::ops::Softmax::Call(
         handle,
         config,
         input_meta.tensor(input),

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) {