[2026春季][T2-1-3]xuzheng567

include/infinicore/ops.hpp

@@ -30,6 +30,7 @@
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"
 #include "ops/fused_gated_delta_net_gating.hpp"
+#include "ops/fused_moe.hpp"
 #include "ops/gelu.hpp"
 #include "ops/gelutanh.hpp"
 #include "ops/hardswish.hpp"

include/infinicore/ops/fused_moe.hpp

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+enum class FusedMoeActivation : int {
+    Silu = 0,
+    Swiglu = 1,
+};
+
+INFINICORE_GRAPH_OP_CLASS(FusedMoe,
+                          Tensor,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          std::optional<Tensor>,
+                          std::optional<Tensor>,
+                          FusedMoeActivation);
+
+Tensor fused_moe(const Tensor &input,
+                 const Tensor &token_selected_experts,
+                 const Tensor &token_final_scales,
+                 const Tensor &w1,
+                 const Tensor &w2,
+                 std::optional<Tensor> b1,
+                 std::optional<Tensor> b2,
+                 FusedMoeActivation activation);
+
+void fused_moe_(Tensor out,
+                const Tensor &input,
+                const Tensor &token_selected_experts,
+                const Tensor &token_final_scales,
+                const Tensor &w1,
+                const Tensor &w2,
+                std::optional<Tensor> b1,
+                std::optional<Tensor> b2,
+                FusedMoeActivation activation);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -57,6 +57,7 @@
 #include "infiniop/ops/fmin.h"
 #include "infiniop/ops/fmod.h"
 #include "infiniop/ops/fused_gated_delta_net_gating.h"
+#include "infiniop/ops/fused_moe.h"
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gelutanh.h"
 #include "infiniop/ops/gemm.h"

include/infiniop/ops/fused_moe.h

@@ -0,0 +1,47 @@
+#ifndef __INFINIOP_FUSED_MOE_API_H__
+#define __INFINIOP_FUSED_MOE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopFusedMoeDescriptor_t;
+
+typedef enum {
+    INFINIOP_FUSED_MOE_ACT_SILU = 0,
+    INFINIOP_FUSED_MOE_ACT_SWIGLU = 1,
+} infiniopFusedMoeActivation_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateFusedMoeDescriptor(
+    infiniopHandle_t handle,
+    infiniopFusedMoeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t token_selected_experts_desc,
+    infiniopTensorDescriptor_t token_final_scales_desc,
+    infiniopTensorDescriptor_t w1_desc,
+    infiniopTensorDescriptor_t w2_desc,
+    infiniopTensorDescriptor_t b1_desc,
+    infiniopTensorDescriptor_t b2_desc,
+    infiniopFusedMoeActivation_t activation);
+
+__INFINI_C __export infiniStatus_t infiniopGetFusedMoeWorkspaceSize(
+    infiniopFusedMoeDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopFusedMoe(
+    infiniopFusedMoeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *input,
+    const void *token_selected_experts,
+    const void *token_final_scales,
+    const void *w1,
+    const void *w2,
+    const void *b1,
+    const void *b2,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyFusedMoeDescriptor(
+    infiniopFusedMoeDescriptor_t desc);
+
+#endif // __INFINIOP_FUSED_MOE_API_H__

python/infinicore/nn/functional/__init__.py

@@ -12,6 +12,7 @@
 from .embedding import embedding
 from .flash_attention import flash_attention
 from .fused_gated_delta_net_gating import fused_gated_delta_net_gating
+from .fused_moe import ACT_SILU, ACT_SWIGLU, fused_moe
 from .gaussian_nll_loss import gaussian_nll_loss
 from .hardswish import hardswish
 from .hardtanh import hardtanh
@@ -54,6 +55,9 @@
     "embedding",
     "flash_attention",
     "fused_gated_delta_net_gating",
+    "fused_moe",
+    "ACT_SILU",
+    "ACT_SWIGLU",
     "gaussian_nll_loss",
     "interpolate",
     "linear",

python/infinicore/nn/functional/fused_moe.py

@@ -0,0 +1,47 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+ACT_SILU = 0
+ACT_SWIGLU = 1
+
+
+def fused_moe(
+    input: Tensor,
+    token_selected_experts: Tensor,
+    token_final_scales: Tensor,
+    w1: Tensor,
+    w2: Tensor,
+    *,
+    b1: Tensor | None = None,
+    b2: Tensor | None = None,
+    activation: int = ACT_SWIGLU,
+    out: Tensor | None = None,
+) -> Tensor:
+    b1_arg = None if b1 is None else b1._underlying
+    b2_arg = None if b2 is None else b2._underlying
+    if out is None:
+        return Tensor(
+            _infinicore.fused_moe(
+                input._underlying,
+                token_selected_experts._underlying,
+                token_final_scales._underlying,
+                w1._underlying,
+                w2._underlying,
+                b1_arg,
+                b2_arg,
+                activation,
+            )
+        )
+
+    _infinicore.fused_moe_(
+        out._underlying,
+        input._underlying,
+        token_selected_experts._underlying,
+        token_final_scales._underlying,
+        w1._underlying,
+        w2._underlying,
+        b1_arg,
+        b2_arg,
+        activation,
+    )
+    return out

src/infinicore/ops/fused_moe/fused_moe.cc

@@ -0,0 +1,66 @@
+#include "infinicore/ops/fused_moe.hpp"
+#include "../../utils.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FusedMoe);
+
+FusedMoe::FusedMoe(Tensor out,
+                   const Tensor &input,
+                   const Tensor &token_selected_experts,
+                   const Tensor &token_final_scales,
+                   const Tensor &w1,
+                   const Tensor &w2,
+                   std::optional<Tensor> b1,
+                   std::optional<Tensor> b2,
+                   FusedMoeActivation activation) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, token_selected_experts, token_final_scales, w1, w2);
+    if (b1.has_value()) {
+        INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, b1.value());
+    }
+    if (b2.has_value()) {
+        INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, b2.value());
+    }
+    INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, token_selected_experts,
+                                 token_final_scales, w1, w2, b1, b2, activation);
+}
+
+void FusedMoe::execute(Tensor out,
+                       const Tensor &input,
+                       const Tensor &token_selected_experts,
+                       const Tensor &token_final_scales,
+                       const Tensor &w1,
+                       const Tensor &w2,
+                       std::optional<Tensor> b1,
+                       std::optional<Tensor> b2,
+                       FusedMoeActivation activation) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(FusedMoe, out, input, token_selected_experts,
+                                      token_final_scales, w1, w2, b1, b2, activation);
+}
+
+Tensor fused_moe(const Tensor &input,
+                 const Tensor &token_selected_experts,
+                 const Tensor &token_final_scales,
+                 const Tensor &w1,
+                 const Tensor &w2,
+                 std::optional<Tensor> b1,
+                 std::optional<Tensor> b2,
+                 FusedMoeActivation activation) {
+    auto out = Tensor::empty(input->shape(), input->dtype(), input->device());
+    fused_moe_(out, input, token_selected_experts, token_final_scales, w1, w2, b1, b2, activation);
+    return out;
+}
+
+void fused_moe_(Tensor out,
+                const Tensor &input,
+                const Tensor &token_selected_experts,
+                const Tensor &token_final_scales,
+                const Tensor &w1,
+                const Tensor &w2,
+                std::optional<Tensor> b1,
+                std::optional<Tensor> b2,
+                FusedMoeActivation activation) {
+    FusedMoe::execute(out, input, token_selected_experts, token_final_scales, w1, w2, b1, b2, activation);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/fused_moe/fused_moe_infiniop.cc

@@ -0,0 +1,80 @@
+#include "infinicore/ops/fused_moe.hpp"
+
+#include "../infiniop_impl.hpp"
+
+namespace infinicore::op::fused_moe_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, FusedMoe, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace, out, input, token_selected_experts, token_final_scales, w1, w2;
+    std::optional<graph::GraphTensor> b1;
+    std::optional<graph::GraphTensor> b2;
+};
+
+void *plan(Tensor out,
+           const Tensor &input,
+           const Tensor &token_selected_experts,
+           const Tensor &token_final_scales,
+           const Tensor &w1,
+           const Tensor &w2,
+           std::optional<Tensor> b1,
+           std::optional<Tensor> b2,
+           FusedMoeActivation activation) {
+    size_t seed = hash_combine(out, input, token_selected_experts, token_final_scales, w1, w2, b1, b2, static_cast<int>(activation));
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, FusedMoe,
+        seed,
+        out->desc(),
+        input->desc(),
+        token_selected_experts->desc(),
+        token_final_scales->desc(),
+        w1->desc(),
+        w2->desc(),
+        b1.has_value() ? b1.value()->desc() : nullptr,
+        b2.has_value() ? b2.value()->desc() : nullptr,
+        static_cast<infiniopFusedMoeActivation_t>(activation));
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, FusedMoe, descriptor);
+
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(out),
+        graph::GraphTensor(input),
+        graph::GraphTensor(token_selected_experts),
+        graph::GraphTensor(token_final_scales),
+        graph::GraphTensor(w1),
+        graph::GraphTensor(w2),
+        b1.has_value() ? std::optional<graph::GraphTensor>(graph::GraphTensor(b1.value())) : std::nullopt,
+        b2.has_value() ? std::optional<graph::GraphTensor>(graph::GraphTensor(b2.value())) : std::nullopt};
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+
+    INFINICORE_CHECK_ERROR(infiniopFusedMoe(
+        planned->descriptor->desc,
+        planned->workspace->data(),
+        planned->workspace->numel(),
+        planned->out->data(),
+        planned->input->data(),
+        planned->token_selected_experts->data(),
+        planned->token_final_scales->data(),
+        planned->w1->data(),
+        planned->w2->data(),
+        planned->b1.has_value() ? planned->b1.value()->data() : nullptr,
+        planned->b2.has_value() ? planned->b2.value()->data() : nullptr,
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(FusedMoe, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::fused_moe_impl::infiniop

src/infinicore/pybind11/ops.hpp

@@ -51,6 +51,7 @@
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"
 #include "ops/fused_gated_delta_net_gating.hpp"
+#include "ops/fused_moe.hpp"
 #include "ops/gaussian_nll_loss.hpp"
 #include "ops/hardswish.hpp"
 #include "ops/hardtanh.hpp"
@@ -172,6 +173,7 @@ inline void bind(py::module &m) {
     bind_kv_caching(m);
     bind_fmod(m);
     bind_fused_gated_delta_net_gating(m);
+    bind_fused_moe(m);
     bind_fmin(m);
     bind_cat(m);
     bind_causal_softmax(m);