Add topkrouter op wrapper

Made-with: Cursor

Author: Ceng23333

include/infinicore/ops.hpp

@@ -38,3 +38,4 @@
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"
 #include "ops/swiglu.hpp"
+#include "ops/topkrouter.hpp"

include/infinicore/ops/topkrouter.hpp

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class TopKRouter {
+public:
+    // values_output: [N, topk] float32
+    // indices_output: [N, topk] int32
+    // input: [N, width] float16/bfloat16/float32
+    // correction_bias: [width] float32
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, size_t);
+    static void execute(Tensor values_output,
+                        Tensor indices_output,
+                        Tensor input,
+                        Tensor correction_bias,
+                        float routed_scaling_factor,
+                        size_t topk);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::pair<Tensor, Tensor> topkrouter(Tensor input,
+                                    Tensor correction_bias,
+                                    float routed_scaling_factor,
+                                    size_t topk);
+
+void topkrouter_(Tensor values_output,
+                 Tensor indices_output,
+                 Tensor input,
+                 Tensor correction_bias,
+                 float routed_scaling_factor,
+                 size_t topk);
+
+} // namespace infinicore::op
+
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class TopKRouter {
+public:
+    // values_output: [N, topk] float32
+    // indices_output: [N, topk] int32
+    // input: [N, width] float16/bfloat16/float32
+    // correction_bias: [width] float32
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, size_t);
+    static void execute(Tensor values_output,
+                        Tensor indices_output,
+                        Tensor input,
+                        Tensor correction_bias,
+                        float routed_scaling_factor,
+                        size_t topk);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::pair<Tensor, Tensor> topkrouter(Tensor input,
+                                    Tensor correction_bias,
+                                    float routed_scaling_factor,
+                                    size_t topk);
+
+void topkrouter_(Tensor values_output,
+                 Tensor indices_output,
+                 Tensor input,
+                 Tensor correction_bias,
+                 float routed_scaling_factor,
+                 size_t topk);
+
+} // namespace infinicore::op
+

src/infinicore/ops/topkrouter/topkrouter.cc

@@ -0,0 +1,55 @@
+#include "infinicore/ops/topkrouter.hpp"
+
+#include "../../utils.hpp"
+
+#include <stdexcept>
+
+namespace infinicore::op {
+
+common::OpDispatcher<TopKRouter::schema> &TopKRouter::dispatcher() {
+    static common::OpDispatcher<TopKRouter::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void TopKRouter::execute(Tensor values_output,
+                         Tensor indices_output,
+                         Tensor input,
+                         Tensor correction_bias,
+                         float routed_scaling_factor,
+                         size_t topk) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values_output, indices_output, input, correction_bias);
+    infinicore::context::setDevice(input->device());
+    auto device_type = input->device().getType();
+    auto func = dispatcher().lookup(device_type);
+    if (func == nullptr) {
+        throw std::runtime_error("No TopKRouter implementation found for device type: " + std::to_string(static_cast<int>(device_type)));
+    }
+    func(values_output, indices_output, input, correction_bias, routed_scaling_factor, topk);
+}
+
+std::pair<Tensor, Tensor> topkrouter(Tensor input,
+                                    Tensor correction_bias,
+                                    float routed_scaling_factor,
+                                    size_t topk) {
+    // values: float32, indices: int32
+    auto shape = input->shape();
+    if (shape.size() != 2) {
+        throw std::runtime_error("topkrouter: input must be 2D [N, width]");
+    }
+    Tensor values = Tensor::empty({shape[0], topk}, DataType::F32, input->device());
+    Tensor indices = Tensor::empty({shape[0], topk}, DataType::I32, input->device());
+    topkrouter_(values, indices, input, correction_bias, routed_scaling_factor, topk);
+    return {values, indices};
+}
+
+void topkrouter_(Tensor values_output,
+                 Tensor indices_output,
+                 Tensor input,
+                 Tensor correction_bias,
+                 float routed_scaling_factor,
+                 size_t topk) {
+    TopKRouter::execute(values_output, indices_output, input, correction_bias, routed_scaling_factor, topk);
+}
+
+} // namespace infinicore::op
+

src/infinicore/ops/topkrouter/topkrouter_infiniop.cc

@@ -0,0 +1,59 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/topkrouter.hpp"
+
+#include <infiniop.h>
+
+namespace infinicore::op::topkrouter_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopTopkrouterDescriptor_t> caches(
+    100,
+    [](infiniopTopkrouterDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyTopkrouterDescriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor values_output,
+               Tensor indices_output,
+               Tensor input,
+               Tensor correction_bias,
+               float routed_scaling_factor,
+               size_t topk) {
+    size_t seed = hash_combine(values_output, indices_output, input, correction_bias, (size_t)topk);
+
+    auto device = context::getDevice();
+    auto &cache = caches.getCache(device);
+
+    auto desc_opt = cache.get(seed);
+    infiniopTopkrouterDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateTopkrouterDescriptor(
+            context::getInfiniopHandle(device), &desc,
+            input->desc(), correction_bias->desc()));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetTopkrouterWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopTopkrouter(
+        desc, workspace->data(), workspace_size,
+        values_output->data(), indices_output->data(),
+        input->data(), correction_bias->data(),
+        routed_scaling_factor, topk, context::getStream()));
+}
+
+static bool registered = []() {
+    TopKRouter::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::topkrouter_impl::infiniop
+