add top_k router and casting

Ceng23333 · Ceng23333 · commit fa7273d076ea · 2026-04-14T03:05:00.000Z
Signed-off-by: Ceng23333 &lt;441651826@qq.com&gt;
diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp
@@ -14,6 +14,7 @@
 #include "ops/binary_cross_entropy_with_logits.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/cdist.hpp"
+#include "ops/convert_to_f32.hpp"
 #include "ops/cross_entropy.hpp"
 #include "ops/embedding.hpp"
 #include "ops/flash_attention.hpp"
diff --git a/include/infinicore/ops/convert_to_f32.hpp b/include/infinicore/ops/convert_to_f32.hpp
@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class ConvertToF32 {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor convert_to_f32(Tensor input);
+void convert_to_f32_(Tensor output, Tensor input);
+
+} // namespace infinicore::op
diff --git a/include/infinicore/ops/topkrouter.hpp b/include/infinicore/ops/topkrouter.hpp
@@ -34,41 +34,3 @@ void topkrouter_(Tensor values_output,
                  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
-
diff --git a/include/infiniop.h b/include/infiniop.h
@@ -26,6 +26,7 @@
 #include "infiniop/ops/cdist.h"
 #include "infiniop/ops/clip.h"
 #include "infiniop/ops/conv.h"
+#include "infiniop/ops/convert_to_f32.h"
 #include "infiniop/ops/cross_entropy.h"
 #include "infiniop/ops/dequant/per_tensor_dequant_int8.h"
 #include "infiniop/ops/dequantize_awq.h"
diff --git a/include/infiniop/ops/convert_to_f32.h b/include/infiniop/ops/convert_to_f32.h
@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_CONVERT_TO_F32_API_H__
+#define __INFINIOP_CONVERT_TO_F32_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopConvertToF32Descriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateConvertToF32Descriptor(infiniopHandle_t handle,
+                                                                          infiniopConvertToF32Descriptor_t *desc_ptr,
+                                                                          infiniopTensorDescriptor_t y,
+                                                                          infiniopTensorDescriptor_t x);
+
+__INFINI_C __export infiniStatus_t infiniopGetConvertToF32WorkspaceSize(infiniopConvertToF32Descriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopConvertToF32(infiniopConvertToF32Descriptor_t desc,
+                                                        void *workspace,
+                                                        size_t workspace_size,
+                                                        void *y,
+                                                        const void *x,
+                                                        void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyConvertToF32Descriptor(infiniopConvertToF32Descriptor_t desc);
+
+#endif
diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py
@@ -69,6 +69,7 @@
 from infinicore.ops.broadcast_to import broadcast_to
 from infinicore.ops.cat import cat
 from infinicore.ops.cdist import cdist
+from infinicore.ops.convert_to_f32 import convert_to_f32
 from infinicore.ops.cross_entropy import cross_entropy
 from infinicore.ops.diff import diff
 from infinicore.ops.digamma import digamma
@@ -216,6 +217,7 @@
     "unsqueeze",
     "rearrange",
     "cross_entropy",
+    "convert_to_f32",
     "tan",
     "empty",
     "empty_like",
diff --git a/python/infinicore/ops/convert_to_f32.py b/python/infinicore/ops/convert_to_f32.py
@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def convert_to_f32(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.convert_to_f32(input._underlying))
+
+    _infinicore.convert_to_f32_(out._underlying, input._underlying)
+
+    return out
diff --git a/src/infinicore/ops/convert_to_f32/convert_to_f32.cc b/src/infinicore/ops/convert_to_f32/convert_to_f32.cc
@@ -0,0 +1,43 @@
+#include "infinicore/ops/convert_to_f32.hpp"
+
+#include "../../utils.hpp"
+#include "infinicore/dtype.hpp"
+
+#include <stdexcept>
+
+namespace infinicore::op {
+
+common::OpDispatcher<ConvertToF32::schema> &ConvertToF32::dispatcher() {
+    static common::OpDispatcher<ConvertToF32::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void ConvertToF32::execute(Tensor output, Tensor input) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input);
+    if (output->dtype() != DataType::F32) {
+        throw std::runtime_error("convert_to_f32: output dtype must be F32");
+    }
+    infinicore::context::setDevice(output->device());
+    auto device_type = output->device().getType();
+    auto func = dispatcher().lookup(device_type);
+
+    if (func == nullptr) {
+        throw std::runtime_error("No convert_to_f32 implementation found for device type: "
+                                 + std::to_string(static_cast<int>(device_type)));
+    }
+
+    func(output, input);
+}
+
+Tensor convert_to_f32(Tensor input) {
+    Shape shape = input->shape();
+    auto output = Tensor::empty(shape, DataType::F32, input->device());
+    convert_to_f32_(output, input);
+    return output;
+}
+
+void convert_to_f32_(Tensor output, Tensor input) {
+    ConvertToF32::execute(output, input);
+}
+
+} // namespace infinicore::op
diff --git a/src/infinicore/ops/convert_to_f32/convert_to_f32_infiniop.cc b/src/infinicore/ops/convert_to_f32/convert_to_f32_infiniop.cc
@@ -0,0 +1,51 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/convert_to_f32.hpp"
+
+#include <infiniop.h>
+
+namespace infinicore::op::convert_to_f32_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopConvertToF32Descriptor_t> caches(
+    100,
+    [](infiniopConvertToF32Descriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyConvertToF32Descriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor output, Tensor input) {
+    size_t seed = hash_combine(output, input);
+
+    auto device = context::getDevice();
+    auto &cache = caches.getCache(device);
+
+    auto desc_opt = cache.get(seed);
+    infiniopConvertToF32Descriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateConvertToF32Descriptor(
+            context::getInfiniopHandle(device), &desc,
+            output->desc(), input->desc()));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetConvertToF32WorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopConvertToF32(
+        desc, workspace->data(), workspace_size,
+        output->data(), input->data(), context::getStream()));
+}
+
+static bool registered = []() {
+    ConvertToF32::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::convert_to_f32_impl::infiniop
diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp
@@ -28,6 +28,7 @@
 #include "ops/cat.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/cdist.hpp"
+#include "ops/convert_to_f32.hpp"
 #include "ops/cross_entropy.hpp"
 #include "ops/diff.hpp"
 #include "ops/digamma.hpp"
@@ -163,6 +164,7 @@ inline void bind(py::module &m) {
     bind_pad(m);
     bind_prelu(m);
     bind_random_sample(m);
+    bind_convert_to_f32(m);
     bind_cross_entropy(m);
     bind_hypot(m);
     bind_take(m);
diff --git a/src/infinicore/pybind11/ops/convert_to_f32.hpp b/src/infinicore/pybind11/ops/convert_to_f32.hpp
@@ -0,0 +1,24 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/convert_to_f32.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_convert_to_f32(py::module &m) {
+    m.def("convert_to_f32",
+          &op::convert_to_f32,
+          py::arg("x"),
+          R"doc(Cast the input tensor to float32 (elementwise).)doc");
+
+    m.def("convert_to_f32_",
+          &op::convert_to_f32_,
+          py::arg("y"),
+          py::arg("x"),
+          R"doc(Cast the input tensor to float32 into ``y`` (elementwise).)doc");
+}
+
+} // namespace infinicore::ops
diff --git a/src/infiniop/ops/convert_to_f32/cpu/convert_to_f32_cpu.cc b/src/infiniop/ops/convert_to_f32/cpu/convert_to_f32_cpu.cc
@@ -0,0 +1,50 @@
+#include "convert_to_f32_cpu.h"
+
+namespace op::convert_to_f32::cpu {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    std::vector<infiniopTensorDescriptor_t> input_desc_vec) {
+
+    auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
+    auto out_dtype = out_desc->dtype();
+    const auto &x_desc = input_desc_vec.at(0);
+    auto in_dtype = x_desc->dtype();
+
+    CHECK_DTYPE(out_dtype, INFINI_DTYPE_F32);
+    CHECK_DTYPE(in_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
+
+    const auto &y_shape = out_desc->shape();
+    const auto &x_shape = x_desc->shape();
+    CHECK_SAME_SHAPE(y_shape, x_shape);
+
+    CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, in_dtype, out_desc, input_desc_vec);
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    std::vector<const void *> inputs,
+    void *stream) const {
+
+    switch (_dtype) {
+    case INFINI_DTYPE_F16:
+        return _device_info->calculate<ConvertToF32Op, float, fp16_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_BF16:
+        return _device_info->calculate<ConvertToF32Op, float, bf16_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_F32:
+        return _device_info->calculate<ConvertToF32Op, float, float>(_info, output, inputs, stream);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace op::convert_to_f32::cpu
diff --git a/src/infiniop/ops/convert_to_f32/cpu/convert_to_f32_cpu.h b/src/infiniop/ops/convert_to_f32/cpu/convert_to_f32_cpu.h
@@ -0,0 +1,19 @@
+#ifndef __CONVERT_TO_F32_CPU_H__
+#define __CONVERT_TO_F32_CPU_H__
+
+#include "../../../elementwise/cpu/elementwise_cpu.h"
+
+ELEMENTWISE_DESCRIPTOR(convert_to_f32, cpu)
+
+namespace op::convert_to_f32::cpu {
+typedef struct ConvertToF32Op {
+public:
+    static constexpr size_t num_inputs = 1;
+    template <typename Tout, typename Tin>
+    Tout operator()(const Tin &x) const {
+        return utils::cast<Tout>(x);
+    }
+} ConvertToF32Op;
+} // namespace op::convert_to_f32::cpu
+
+#endif // __CONVERT_TO_F32_CPU_H__
diff --git a/src/infiniop/ops/convert_to_f32/cuda/kernel.cuh b/src/infiniop/ops/convert_to_f32/cuda/kernel.cuh
@@ -0,0 +1,29 @@
+#ifndef __CONVERT_TO_F32_CUDA_H__
+#define __CONVERT_TO_F32_CUDA_H__
+
+#include <type_traits>
+
+namespace op::convert_to_f32::cuda {
+
+struct ConvertToF32Op {
+public:
+    static constexpr size_t num_inputs = 1;
+
+    template <typename Tout, typename Tin>
+    __device__ __forceinline__ Tout operator()(const Tin &x) const {
+        static_assert(std::is_same_v<Tout, float>, "convert_to_f32 output must be float");
+        if constexpr (std::is_same_v<Tin, half>) {
+            return __half2float(x);
+        } else if constexpr (std::is_same_v<Tin, cuda_bfloat16>) {
+            return __bfloat162float(x);
+        } else if constexpr (std::is_same_v<Tin, float>) {
+            return x;
+        } else {
+            return static_cast<float>(x);
+        }
+    }
+};
+
+} // namespace op::convert_to_f32::cuda
+
+#endif // __CONVERT_TO_F32_CUDA_H__
diff --git a/src/infiniop/ops/convert_to_f32/nvidia/convert_to_f32_nvidia.cu b/src/infiniop/ops/convert_to_f32/nvidia/convert_to_f32_nvidia.cu
diff --git a/src/infiniop/ops/convert_to_f32/nvidia/convert_to_f32_nvidia.cuh b/src/infiniop/ops/convert_to_f32/nvidia/convert_to_f32_nvidia.cuh
diff --git a/src/infiniop/ops/convert_to_f32/operator.cc b/src/infiniop/ops/convert_to_f32/operator.cc
diff --git a/test/infinicore/ops/convert_to_f32.py b/test/infinicore/ops/convert_to_f32.py
diff --git a/test/infiniop/convert_to_f32.py b/test/infiniop/convert_to_f32.py
diff --git a/test/infiniop/libinfiniop/op_register.py b/test/infiniop/libinfiniop/op_register.py
diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua
