From 5736c8c9f936dd0a3955e59f07101dd8e3aaa8a3 Mon Sep 17 00:00:00 2001 From: xuzheng567 <17610366386@163.com> Date: Fri, 10 Jul 2026 16:50:40 +0800 Subject: [PATCH 1/3] add fused moe op --- include/infinicore/ops.hpp | 1 + include/infinicore/ops/fused_moe.hpp | 46 ++++++ include/infiniop.h | 1 + include/infiniop/ops/fused_moe.h | 47 ++++++ python/infinicore/nn/functional/__init__.py | 4 + python/infinicore/nn/functional/fused_moe.py | 47 ++++++ src/infinicore/ops/fused_moe/fused_moe.cc | 66 ++++++++ .../ops/fused_moe/fused_moe_infiniop.cc | 80 +++++++++ src/infinicore/pybind11/ops.hpp | 2 + src/infinicore/pybind11/ops/fused_moe.hpp | 72 ++++++++ src/infiniop/ops/fused_moe/cuda/kernel.cuh | 115 +++++++++++++ src/infiniop/ops/fused_moe/fused_moe.h | 45 +++++ src/infiniop/ops/fused_moe/info.h | 123 ++++++++++++++ .../ops/fused_moe/nvidia/fused_moe_nvidia.cu | 94 +++++++++++ .../ops/fused_moe/nvidia/fused_moe_nvidia.cuh | 8 + src/infiniop/ops/fused_moe/operator.cc | 96 +++++++++++ test/infinicore/ops/fused_moe.py | 99 +++++++++++ test/infiniop/fused_moe.py | 155 ++++++++++++++++++ test/infiniop/libinfiniop/op_register.py | 38 +++++ 19 files changed, 1139 insertions(+) create mode 100644 include/infinicore/ops/fused_moe.hpp create mode 100644 include/infiniop/ops/fused_moe.h create mode 100644 python/infinicore/nn/functional/fused_moe.py create mode 100644 src/infinicore/ops/fused_moe/fused_moe.cc create mode 100644 src/infinicore/ops/fused_moe/fused_moe_infiniop.cc create mode 100644 src/infinicore/pybind11/ops/fused_moe.hpp create mode 100644 src/infiniop/ops/fused_moe/cuda/kernel.cuh create mode 100644 src/infiniop/ops/fused_moe/fused_moe.h create mode 100644 src/infiniop/ops/fused_moe/info.h create mode 100644 src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu create mode 100644 src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cuh create mode 100644 src/infiniop/ops/fused_moe/operator.cc create mode 100644 test/infinicore/ops/fused_moe.py create mode 100644 test/infiniop/fused_moe.py diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b5c4ff18f..ed3abe502 100644 --- a/include/infinicore/ops.hpp +++ b/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" diff --git a/include/infinicore/ops/fused_moe.hpp b/include/infinicore/ops/fused_moe.hpp new file mode 100644 index 000000000..8e153e6f1 --- /dev/null +++ b/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 + +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, + std::optional, + 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 b1, + std::optional 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 b1, + std::optional b2, + FusedMoeActivation activation); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 95bf75a0d..67c8f9bb1 100644 --- a/include/infiniop.h +++ b/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" diff --git a/include/infiniop/ops/fused_moe.h b/include/infiniop/ops/fused_moe.h new file mode 100644 index 000000000..8b022c1c3 --- /dev/null +++ b/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__ diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py index 0395c5626..0569fa1eb 100644 --- a/python/infinicore/nn/functional/__init__.py +++ b/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", diff --git a/python/infinicore/nn/functional/fused_moe.py b/python/infinicore/nn/functional/fused_moe.py new file mode 100644 index 000000000..f301c5e2d --- /dev/null +++ b/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 diff --git a/src/infinicore/ops/fused_moe/fused_moe.cc b/src/infinicore/ops/fused_moe/fused_moe.cc new file mode 100644 index 000000000..978d9285a --- /dev/null +++ b/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 b1, + std::optional 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 b1, + std::optional 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 b1, + std::optional 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 b1, + std::optional b2, + FusedMoeActivation activation) { + FusedMoe::execute(out, input, token_selected_experts, token_final_scales, w1, w2, b1, b2, activation); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fused_moe/fused_moe_infiniop.cc b/src/infinicore/ops/fused_moe/fused_moe_infiniop.cc new file mode 100644 index 000000000..d91bbb2fd --- /dev/null +++ b/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; + graph::GraphTensor workspace, out, input, token_selected_experts, token_final_scales, w1, w2; + std::optional b1; + std::optional 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 b1, + std::optional b2, + FusedMoeActivation activation) { + size_t seed = hash_combine(out, input, token_selected_experts, token_final_scales, w1, w2, b1, b2, static_cast(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(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(b1.value())) : std::nullopt, + b2.has_value() ? std::optional(graph::GraphTensor(b2.value())) : std::nullopt}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(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(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(FusedMoe, &plan, &run, &cleanup); + +} // namespace infinicore::op::fused_moe_impl::infiniop diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 087261382..88dbbcec2 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/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); diff --git a/src/infinicore/pybind11/ops/fused_moe.hpp b/src/infinicore/pybind11/ops/fused_moe.hpp new file mode 100644 index 000000000..e52de0a9a --- /dev/null +++ b/src/infinicore/pybind11/ops/fused_moe.hpp @@ -0,0 +1,72 @@ +#pragma once + +#include + +#include "infinicore/ops/fused_moe.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline std::optional optional_tensor(py::object obj) { + if (obj.is_none()) { + return std::nullopt; + } + return obj.cast(); +} + +inline Tensor py_fused_moe(Tensor input, + Tensor token_selected_experts, + Tensor token_final_scales, + Tensor w1, + Tensor w2, + py::object b1, + py::object b2, + int activation) { + return op::fused_moe(input, token_selected_experts, token_final_scales, w1, w2, + optional_tensor(b1), optional_tensor(b2), + static_cast(activation)); +} + +inline void py_fused_moe_(Tensor out, + Tensor input, + Tensor token_selected_experts, + Tensor token_final_scales, + Tensor w1, + Tensor w2, + py::object b1, + py::object b2, + int activation) { + op::fused_moe_(out, input, token_selected_experts, token_final_scales, w1, w2, + optional_tensor(b1), optional_tensor(b2), + static_cast(activation)); +} + +inline void bind_fused_moe(py::module &m) { + m.def("fused_moe", + &ops::py_fused_moe, + py::arg("input"), + py::arg("token_selected_experts"), + py::arg("token_final_scales"), + py::arg("w1"), + py::arg("w2"), + py::arg("b1") = py::none(), + py::arg("b2") = py::none(), + py::arg("activation") = 1, + R"doc(Fused MoE consuming topksoftmax values/indices. activation: 0=silu, 1=swiglu.)doc"); + + m.def("fused_moe_", + &ops::py_fused_moe_, + py::arg("out"), + py::arg("input"), + py::arg("token_selected_experts"), + py::arg("token_final_scales"), + py::arg("w1"), + py::arg("w2"), + py::arg("b1") = py::none(), + py::arg("b2") = py::none(), + py::arg("activation") = 1, + R"doc(In-place fused MoE consuming topksoftmax values/indices.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infiniop/ops/fused_moe/cuda/kernel.cuh b/src/infiniop/ops/fused_moe/cuda/kernel.cuh new file mode 100644 index 000000000..c1394239e --- /dev/null +++ b/src/infiniop/ops/fused_moe/cuda/kernel.cuh @@ -0,0 +1,115 @@ +#ifndef __FUSED_MOE_KERNEL_CUH__ +#define __FUSED_MOE_KERNEL_CUH__ + +#include "infiniop/ops/fused_moe.h" +#include +#include +#include + +// Inspired by TensorRT-LLM fused MoE post-routing contract: +// cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h and +// cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu. +// TensorRT-LLM is licensed under Apache-2.0. This file is an InfiniCore +// implementation and intentionally does not depend on TensorRT-LLM symbols. + +template +__device__ inline float moeToFloat(T v) { + return static_cast(v); +} + +template <> +__device__ inline float moeToFloat(half v) { + return __half2float(v); +} + +template <> +__device__ inline float moeToFloat<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ inline T moeFromFloat(float v) { + return static_cast(v); +} + +template <> +__device__ inline half moeFromFloat(float v) { + return __float2half(v); +} + +template <> +__device__ inline __nv_bfloat16 moeFromFloat<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} + +__device__ inline float moeSilu(float x) { + return x / (1.0f + expf(-x)); +} + +template +__global__ void fusedMoeKernel( + T *out, + const T *input, + const int32_t *selected_experts, + const float *final_scales, + const T *w1, + const T *w2, + const T *b1, + const T *b2, + size_t N, + size_t hidden_size, + size_t inter_size, + size_t num_experts, + size_t topk, + size_t w1_cols, + int activation) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = N * hidden_size; + if (idx >= total) { + return; + } + + size_t token = idx / hidden_size; + size_t out_h = idx % hidden_size; + float accum = 0.0f; + + for (size_t route = 0; route < topk; ++route) { + int expert = selected_experts[token * topk + route]; + if (expert < 0 || static_cast(expert) >= num_experts) { + continue; + } + float route_scale = final_scales[token * topk + route]; + float expert_out = 0.0f; + for (size_t j = 0; j < inter_size; ++j) { + float gate = 0.0f; + float up = 0.0f; + size_t gate_row = j; + size_t up_row = activation == INFINIOP_FUSED_MOE_ACT_SWIGLU ? (j + inter_size) : j; + const T *w1_expert = w1 + (static_cast(expert) * w1_cols * hidden_size); + for (size_t h = 0; h < hidden_size; ++h) { + float x = moeToFloat(input[token * hidden_size + h]); + gate += moeToFloat(w1_expert[gate_row * hidden_size + h]) * x; + if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { + up += moeToFloat(w1_expert[up_row * hidden_size + h]) * x; + } + } + if (b1 != nullptr) { + const T *b1_expert = b1 + static_cast(expert) * w1_cols; + gate += moeToFloat(b1_expert[gate_row]); + if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { + up += moeToFloat(b1_expert[up_row]); + } + } + float act = activation == INFINIOP_FUSED_MOE_ACT_SWIGLU ? (moeSilu(gate) * up) : moeSilu(gate); + const T *w2_expert = w2 + (static_cast(expert) * hidden_size * inter_size); + expert_out += moeToFloat(w2_expert[out_h * inter_size + j]) * act; + } + if (b2 != nullptr) { + expert_out += moeToFloat(b2[static_cast(expert) * hidden_size + out_h]); + } + accum += route_scale * expert_out; + } + out[idx] = moeFromFloat(accum); +} + +#endif // __FUSED_MOE_KERNEL_CUH__ diff --git a/src/infiniop/ops/fused_moe/fused_moe.h b/src/infiniop/ops/fused_moe/fused_moe.h new file mode 100644 index 000000000..a838b6106 --- /dev/null +++ b/src/infiniop/ops/fused_moe/fused_moe.h @@ -0,0 +1,45 @@ +#ifndef __INFINIOP_FUSED_MOE_H__ +#define __INFINIOP_FUSED_MOE_H__ + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::fused_moe::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + FusedMoeInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor(Opaque *opaque, FusedMoeInfo info, size_t workspace_size,\ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info), _workspace_size(workspace_size) {}\ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const { return _workspace_size; } \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **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); \ + infiniStatus_t calculate( \ + 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) const; \ + }; \ + } + +#endif // __INFINIOP_FUSED_MOE_H__ diff --git a/src/infiniop/ops/fused_moe/info.h b/src/infiniop/ops/fused_moe/info.h new file mode 100644 index 000000000..089e3e1c5 --- /dev/null +++ b/src/infiniop/ops/fused_moe/info.h @@ -0,0 +1,123 @@ +#ifndef __FUSED_MOE_INFO_H__ +#define __FUSED_MOE_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" +#include "infiniop/ops/fused_moe.h" + +namespace op::fused_moe { + +class FusedMoeInfo { + FusedMoeInfo() = default; + +public: + infiniDtype_t dtype; + bool has_b1; + bool has_b2; + infiniopFusedMoeActivation_t activation; + size_t N; + size_t hidden_size; + size_t inter_size; + size_t num_experts; + size_t topk; + size_t w1_cols; + + static utils::Result create( + 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) { + + if (out_desc == nullptr || input_desc == nullptr || + token_selected_experts_desc == nullptr || token_final_scales_desc == nullptr || + w1_desc == nullptr || w2_desc == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (activation != INFINIOP_FUSED_MOE_ACT_SILU && activation != INFINIOP_FUSED_MOE_ACT_SWIGLU) { + return INFINI_STATUS_BAD_PARAM; + } + + auto dtype = input_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + + if (out_desc->dtype() != dtype || w1_desc->dtype() != dtype || w2_desc->dtype() != dtype || + (b1_desc != nullptr && b1_desc->dtype() != dtype) || + (b2_desc != nullptr && b2_desc->dtype() != dtype)) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (token_selected_experts_desc->dtype() != INFINI_DTYPE_I32 || + token_final_scales_desc->dtype() != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + if (input_desc->ndim() != 2 || out_desc->ndim() != 2 || + token_selected_experts_desc->ndim() != 2 || token_final_scales_desc->ndim() != 2 || + w1_desc->ndim() != 3 || w2_desc->ndim() != 3 || + (b1_desc != nullptr && b1_desc->ndim() != 2) || + (b2_desc != nullptr && b2_desc->ndim() != 2)) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + auto input_shape = input_desc->shape(); + auto out_shape = out_desc->shape(); + auto indices_shape = token_selected_experts_desc->shape(); + auto scales_shape = token_final_scales_desc->shape(); + auto w1_shape = w1_desc->shape(); + auto w2_shape = w2_desc->shape(); + + size_t N = input_shape[0]; + size_t hidden = input_shape[1]; + size_t topk = indices_shape[1]; + size_t experts = w1_shape[0]; + size_t w1_cols = w1_shape[1]; + size_t inter = w2_shape[2]; + + if (out_shape[0] != N || out_shape[1] != hidden || + indices_shape[0] != N || scales_shape[0] != N || scales_shape[1] != topk || + w1_shape[2] != hidden || w2_shape[0] != experts || w2_shape[1] != hidden) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { + if (w1_cols != 2 * inter) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } else if (w1_cols != inter) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + if (b1_desc != nullptr) { + auto b1_shape = b1_desc->shape(); + if (b1_shape[0] != experts || b1_shape[1] != w1_cols) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + if (b2_desc != nullptr) { + auto b2_shape = b2_desc->shape(); + if (b2_shape[0] != experts || b2_shape[1] != hidden) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + + FusedMoeInfo info; + info.dtype = dtype; + info.has_b1 = b1_desc != nullptr; + info.has_b2 = b2_desc != nullptr; + info.activation = activation; + info.N = N; + info.hidden_size = hidden; + info.inter_size = inter; + info.num_experts = experts; + info.topk = topk; + info.w1_cols = w1_cols; + return utils::Result(info); + } +}; + +} // namespace op::fused_moe + +#endif // __FUSED_MOE_INFO_H__ diff --git a/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu new file mode 100644 index 000000000..edc215283 --- /dev/null +++ b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu @@ -0,0 +1,94 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "fused_moe_nvidia.cuh" + +#include "../cuda/kernel.cuh" +#include + +namespace op::fused_moe::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **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) { + auto info = FusedMoeInfo::create(out_desc, input_desc, token_selected_experts_desc, + token_final_scales_desc, w1_desc, w2_desc, + b1_desc, b2_desc, activation); + CHECK_RESULT(info); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + 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_) const { + if (out == nullptr || input == nullptr || token_selected_experts == nullptr || + token_final_scales == nullptr || w1 == nullptr || w2 == nullptr || + (_info.has_b1 && b1 == nullptr) || (_info.has_b2 && b2 == nullptr)) { + return INFINI_STATUS_NULL_POINTER; + } + + cudaStream_t stream = (cudaStream_t)stream_; + size_t total = _info.N * _info.hidden_size; + int threads = 256; + int blocks = static_cast((total + threads - 1) / threads); + + if (_info.dtype == INFINI_DTYPE_F16) { + fusedMoeKernel<<>>( + static_cast(out), static_cast(input), + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w1), static_cast(w2), + static_cast(b1), static_cast(b2), + _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, + _info.topk, _info.w1_cols, static_cast(_info.activation)); + } else if (_info.dtype == INFINI_DTYPE_BF16) { + fusedMoeKernel<__nv_bfloat16><<>>( + static_cast<__nv_bfloat16 *>(out), static_cast(input), + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w1), static_cast(w2), + static_cast(b1), static_cast(b2), + _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, + _info.topk, _info.w1_cols, static_cast(_info.activation)); + } else if (_info.dtype == INFINI_DTYPE_F32) { + fusedMoeKernel<<>>( + static_cast(out), static_cast(input), + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w1), static_cast(w2), + static_cast(b1), static_cast(b2), + _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, + _info.topk, _info.w1_cols, static_cast(_info.activation)); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::fused_moe::nvidia diff --git a/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cuh b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cuh new file mode 100644 index 000000000..bc4a9597d --- /dev/null +++ b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FUSED_MOE_NVIDIA_CUH__ +#define __FUSED_MOE_NVIDIA_CUH__ + +#include "../fused_moe.h" + +DESCRIPTOR(nvidia) + +#endif // __FUSED_MOE_NVIDIA_CUH__ diff --git a/src/infiniop/ops/fused_moe/operator.cc b/src/infiniop/ops/fused_moe/operator.cc new file mode 100644 index 000000000..b7e004714 --- /dev/null +++ b/src/infiniop/ops/fused_moe/operator.cc @@ -0,0 +1,96 @@ +#include "../../handle.h" +#include "../../operator.h" +#include "infiniop/ops/fused_moe.h" + +#ifdef ENABLE_NVIDIA_API +#include "nvidia/fused_moe_nvidia.cuh" +#endif + +__INFINI_C 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) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fused_moe::NAMESPACE::Descriptor::create( \ + handle, reinterpret_cast(desc_ptr), \ + out_desc, input_desc, token_selected_experts_desc, token_final_scales_desc, \ + w1_desc, w2_desc, b1_desc, b2_desc, activation); + + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetFusedMoeWorkspaceSize(infiniopFusedMoeDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C 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) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate( \ + workspace, workspace_size, out, input, token_selected_experts, token_final_scales, \ + w1, w2, b1, b2, stream); + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFusedMoeDescriptor(infiniopFusedMoeDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/test/infinicore/ops/fused_moe.py b/test/infinicore/ops/fused_moe.py new file mode 100644 index 000000000..202db27c8 --- /dev/null +++ b/test/infinicore/ops/fused_moe.py @@ -0,0 +1,99 @@ +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +import torch +import torch.nn.functional as F +import infinicore +from framework import get_args, get_test_devices, torch_device_map, InfiniDeviceEnum, to_torch_dtype, convert_infinicore_to_torch + +ACT_SILU = 0 +ACT_SWIGLU = 1 +CASES = [ + (2, 16, 32, 4, 2, ACT_SILU), + (3, 32, 16, 5, 2, ACT_SWIGLU), +] +DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] +TOLS = { + infinicore.float16: {"atol": 2e-2, "rtol": 2e-2}, + infinicore.bfloat16: {"atol": 5e-2, "rtol": 5e-2}, + infinicore.float32: {"atol": 1e-4, "rtol": 1e-4}, +} + + +def ref(x, indices, scales, w1, w2, b1, b2, activation): + N, hidden = x.shape + topk = indices.shape[1] + out = torch.zeros((N, hidden), dtype=torch.float32, device=x.device) + for n in range(N): + x_f = x[n].float() + for k in range(topk): + e = int(indices[n, k]) + h1 = w1[e].float() @ x_f + if b1 is not None: + h1 = h1 + b1[e].float() + if activation == ACT_SWIGLU: + gate, up = h1.chunk(2, dim=0) + act = F.silu(gate) * up + else: + act = F.silu(h1) + y = w2[e].float() @ act + if b2 is not None: + y = y + b2[e].float() + out[n] += scales[n, k].float() * y + return out.to(x.dtype) + + +def wrap(t): + return infinicore.from_torch(t.contiguous()) + + +def run_case(device, case, dtype): + N, hidden, inter, experts, topk, activation = case + torch_device = torch_device_map[device] + torch_dtype = to_torch_dtype(dtype) + print(f"Testing InfiniCore fused_moe N={N} hidden={hidden} inter={inter} experts={experts} topk={topk} activation={activation} dtype={dtype}") + x = (torch.rand((N, hidden), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + w1_cols = inter * 2 if activation == ACT_SWIGLU else inter + w1 = (torch.rand((experts, w1_cols, hidden), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + w2 = (torch.rand((experts, hidden, inter), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + b1 = (torch.rand((experts, w1_cols), dtype=torch_dtype, device=torch_device) * 0.1).contiguous() + b2 = (torch.rand((experts, hidden), dtype=torch_dtype, device=torch_device) * 0.1).contiguous() + logits = torch.rand((N, experts), dtype=torch.float32, device=torch_device) + scales, indices64 = torch.topk(F.softmax(logits, dim=-1), topk, dim=-1) + scales = (scales / scales.sum(dim=-1, keepdim=True)).contiguous() + indices = indices64.to(torch.int32).contiguous() + ans = ref(x, indices, scales, w1, w2, b1, b2, activation) + + out = infinicore.nn.functional.fused_moe( + wrap(x), wrap(indices), wrap(scales), wrap(w1), wrap(w2), b1=wrap(b1), b2=wrap(b2), activation=activation + ) + infinicore.sync_device() + actual = convert_infinicore_to_torch(out) + assert torch.allclose(actual, ans, **TOLS[dtype]) + + out_inplace = infinicore.empty((N, hidden), dtype=dtype, device=infinicore.device(torch_device, 0)) + returned = infinicore.nn.functional.fused_moe( + wrap(x), wrap(indices), wrap(scales), wrap(w1), wrap(w2), b1=wrap(b1), b2=wrap(b2), activation=activation, out=out_inplace + ) + infinicore.sync_device() + assert returned is out_inplace + actual_inplace = convert_infinicore_to_torch(out_inplace) + assert torch.allclose(actual_inplace, ans, **TOLS[dtype]) + + +def main(): + args = get_args() + for device in get_test_devices(args): + if device != InfiniDeviceEnum.NVIDIA: + continue + infinicore.set_device(infinicore.device(torch_device_map[device], 0)) + for case in CASES: + for dtype in DTYPES: + run_case(device, case, dtype) + print("\033[92mTest passed!\033[0m") + + +if __name__ == "__main__": + main() diff --git a/test/infiniop/fused_moe.py b/test/infiniop/fused_moe.py new file mode 100644 index 000000000..f00c1bceb --- /dev/null +++ b/test/infiniop/fused_moe.py @@ -0,0 +1,155 @@ +import ctypes +from ctypes import c_uint64, c_int32 +import torch +import torch.nn.functional as F + +from libinfiniop import ( + LIBINFINIOP, + TestTensor, + get_test_devices, + check_error, + test_operator, + get_args, + debug, + get_tolerance, + TestWorkspace, + InfiniDtype, + InfiniDtypeNames, + InfiniDeviceNames, + infiniopOperatorDescriptor_t, + torch_device_map, +) + +ACT_SILU = 0 +ACT_SWIGLU = 1 + +_TEST_CASES_ = [ + # N, hidden, inter, experts, topk, activation + (2, 16, 32, 4, 2, ACT_SILU), + (3, 32, 16, 5, 2, ACT_SWIGLU), + (1, 16, 16, 3, 1, ACT_SWIGLU), +] +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16, InfiniDtype.F32] +_TOLERANCE_MAP = { + InfiniDtype.F16: {"atol": 2e-2, "rtol": 2e-2}, + InfiniDtype.BF16: {"atol": 5e-2, "rtol": 5e-2}, + InfiniDtype.F32: {"atol": 1e-4, "rtol": 1e-4}, +} + +DEBUG = False + + +def torch_fused_moe(x, indices, scales, w1, w2, b1, b2, activation): + N, hidden = x.shape + topk = indices.shape[1] + out = torch.zeros((N, hidden), dtype=torch.float32, device=x.device) + for n in range(N): + x_f = x[n].float() + for k in range(topk): + expert = int(indices[n, k].item()) + hidden1 = torch.matmul(w1[expert].float(), x_f) + if b1 is not None: + hidden1 = hidden1 + b1[expert].float() + if activation == ACT_SWIGLU: + gate, up = hidden1.chunk(2, dim=0) + act = F.silu(gate) * up + else: + act = F.silu(hidden1) + y = torch.matmul(w2[expert].float(), act) + if b2 is not None: + y = y + b2[expert].float() + out[n] += scales[n, k].float() * y + return out.to(x.dtype) + + +def test(handle, device, N, hidden, inter, experts, topk, activation, dtype=InfiniDtype.F16, sync=None): + print( + f"Testing FusedMoe on {InfiniDeviceNames[device]} N={N} hidden={hidden} inter={inter} " + f"experts={experts} topk={topk} activation={activation} dtype={InfiniDtypeNames[dtype]}" + ) + torch_device = torch_device_map[device] + torch_dtype = {InfiniDtype.F16: torch.float16, InfiniDtype.BF16: torch.bfloat16, InfiniDtype.F32: torch.float32}[dtype] + + x_t = (torch.rand((N, hidden), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + w1_cols = inter * 2 if activation == ACT_SWIGLU else inter + w1_t = (torch.rand((experts, w1_cols, hidden), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + w2_t = (torch.rand((experts, hidden, inter), dtype=torch_dtype, device=torch_device) * 2 - 1).contiguous() + b1_t = (torch.rand((experts, w1_cols), dtype=torch_dtype, device=torch_device) * 0.1).contiguous() + b2_t = (torch.rand((experts, hidden), dtype=torch_dtype, device=torch_device) * 0.1).contiguous() + + logits = torch.rand((N, experts), dtype=torch.float32, device=torch_device) + scales_t, indices_i64 = torch.topk(F.softmax(logits, dim=-1), topk, dim=-1) + scales_t = scales_t / scales_t.sum(dim=-1, keepdim=True) + indices_t = indices_i64.to(torch.int32).contiguous() + scales_t = scales_t.contiguous() + + ans = torch_fused_moe(x_t, indices_t, scales_t, w1_t, w2_t, b1_t, b2_t, activation) + + x = TestTensor.from_torch(x_t, dtype, device) + w1 = TestTensor.from_torch(w1_t, dtype, device) + w2 = TestTensor.from_torch(w2_t, dtype, device) + b1 = TestTensor.from_torch(b1_t, dtype, device) + b2 = TestTensor.from_torch(b2_t, dtype, device) + indices = TestTensor.from_torch(indices_t, InfiniDtype.I32, device) + scales = TestTensor.from_torch(scales_t, InfiniDtype.F32, device) + out = TestTensor((N, hidden), None, dtype, device) + + if sync: + sync() + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFusedMoeDescriptor( + handle, + ctypes.byref(descriptor), + out.descriptor, + x.descriptor, + indices.descriptor, + scales.descriptor, + w1.descriptor, + w2.descriptor, + b1.descriptor, + b2.descriptor, + c_int32(activation), + ) + ) + + workspace_size = c_uint64(0) + check_error(LIBINFINIOP.infiniopGetFusedMoeWorkspaceSize(descriptor, ctypes.byref(workspace_size))) + workspace = TestWorkspace(workspace_size.value, x.device) + + for tensor in [x, w1, w2, b1, b2, indices, scales, out]: + tensor.destroy_desc() + + check_error( + LIBINFINIOP.infiniopFusedMoe( + descriptor, + workspace.data(), + workspace_size.value, + out.data(), + x.data(), + indices.data(), + scales.data(), + w1.data(), + w2.data(), + b1.data(), + b2.data(), + None, + ) + ) + if sync: + sync() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + if DEBUG: + debug(out.actual_tensor(), ans, atol=atol, rtol=rtol) + assert torch.allclose(out.actual_tensor(), ans, atol=atol, rtol=rtol) + check_error(LIBINFINIOP.infiniopDestroyFusedMoeDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + DEBUG = args.debug + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES_, _TENSOR_DTYPES) + print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/libinfiniop/op_register.py b/test/infiniop/libinfiniop/op_register.py index a45241b56..12ae19e65 100644 --- a/test/infiniop/libinfiniop/op_register.py +++ b/test/infiniop/libinfiniop/op_register.py @@ -132,6 +132,44 @@ def addcmul_(lib): infiniopOperatorDescriptor_t, # descriptor ] +@OpRegister.operator +def fused_moe_(lib): + lib.infiniopCreateFusedMoeDescriptor.restype = c_int32 + lib.infiniopCreateFusedMoeDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_void_p, + c_void_p, + c_int32, + ] + lib.infiniopGetFusedMoeWorkspaceSize.restype = c_int32 + lib.infiniopGetFusedMoeWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopFusedMoe.restype = c_int32 + lib.infiniopFusedMoe.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyFusedMoeDescriptor.restype = c_int32 + lib.infiniopDestroyFusedMoeDescriptor.argtypes = [infiniopOperatorDescriptor_t] @OpRegister.operator def fused_gated_delta_net_gating_(lib): From 4f6c671f4954cedc0e95b8fc58f18fa879d452ea Mon Sep 17 00:00:00 2001 From: xuzheng567 <17610366386@163.com> Date: Fri, 10 Jul 2026 16:50:41 +0800 Subject: [PATCH 2/3] optimize with cutlass --- src/infiniop/ops/fused_moe/cuda/kernel.cuh | 159 +++-- .../ops/fused_moe/nvidia/fused_moe_nvidia.cu | 658 +++++++++++++++++- test/infinicore/ops/fused_moe.py | 2 + 3 files changed, 727 insertions(+), 92 deletions(-) diff --git a/src/infiniop/ops/fused_moe/cuda/kernel.cuh b/src/infiniop/ops/fused_moe/cuda/kernel.cuh index c1394239e..60bf63ede 100644 --- a/src/infiniop/ops/fused_moe/cuda/kernel.cuh +++ b/src/infiniop/ops/fused_moe/cuda/kernel.cuh @@ -1,16 +1,8 @@ #ifndef __FUSED_MOE_KERNEL_CUH__ #define __FUSED_MOE_KERNEL_CUH__ -#include "infiniop/ops/fused_moe.h" -#include -#include #include - -// Inspired by TensorRT-LLM fused MoE post-routing contract: -// cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h and -// cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu. -// TensorRT-LLM is licensed under Apache-2.0. This file is an InfiniCore -// implementation and intentionally does not depend on TensorRT-LLM symbols. +#include template __device__ inline float moeToFloat(T v) { @@ -23,7 +15,7 @@ __device__ inline float moeToFloat(half v) { } template <> -__device__ inline float moeToFloat<__nv_bfloat16>(__nv_bfloat16 v) { +__device__ inline float moeToFloat(cuda_bfloat16 v) { return __bfloat162float(v); } @@ -38,7 +30,7 @@ __device__ inline half moeFromFloat(float v) { } template <> -__device__ inline __nv_bfloat16 moeFromFloat<__nv_bfloat16>(float v) { +__device__ inline cuda_bfloat16 moeFromFloat(float v) { return __float2bfloat16(v); } @@ -47,69 +39,118 @@ __device__ inline float moeSilu(float x) { } template -__global__ void fusedMoeKernel( - T *out, +__global__ void fusedMoeW1Kernel( + T *w1_out, const T *input, const int32_t *selected_experts, - const float *final_scales, const T *w1, - const T *w2, const T *b1, - const T *b2, - size_t N, + size_t route_count, size_t hidden_size, - size_t inter_size, - size_t num_experts, size_t topk, size_t w1_cols, + size_t num_experts) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = route_count * w1_cols; + if (idx >= total) { + return; + } + + size_t route_id = idx / w1_cols; + size_t col = idx % w1_cols; + size_t token = route_id / topk; + int expert = selected_experts[route_id]; + if (expert < 0 || static_cast(expert) >= num_experts) { + w1_out[idx] = moeFromFloat(0.0f); + return; + } + + const T *x = input + token * hidden_size; + const T *w = w1 + (static_cast(expert) * w1_cols + col) * hidden_size; + float acc = 0.0f; + for (size_t h = 0; h < hidden_size; ++h) { + acc += moeToFloat(w[h]) * moeToFloat(x[h]); + } + if (b1 != nullptr) { + acc += moeToFloat(b1[static_cast(expert) * w1_cols + col]); + } + w1_out[idx] = moeFromFloat(acc); +} + +template +__global__ void fusedMoeActivationKernel( + T *activated, + const T *w1_out, + size_t route_count, + size_t inter_size, + size_t w1_cols, int activation) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t total = N * hidden_size; + size_t total = route_count * inter_size; if (idx >= total) { return; } - size_t token = idx / hidden_size; + size_t route_id = idx / inter_size; + size_t j = idx % inter_size; + const T *row = w1_out + route_id * w1_cols; + float gate = moeToFloat(row[j]); + float act; + if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { + float up = moeToFloat(row[j + inter_size]); + act = moeSilu(gate) * up; + } else { + act = moeSilu(gate); + } + activated[idx] = moeFromFloat(act); +} + +template +__global__ void fusedMoeW2ScatterKernel( + float *out_accum, + const T *activated, + const int32_t *selected_experts, + const float *final_scales, + const T *w2, + const T *b2, + size_t route_count, + size_t hidden_size, + size_t inter_size, + size_t topk, + size_t num_experts) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = route_count * hidden_size; + if (idx >= total) { + return; + } + + size_t route_id = idx / hidden_size; size_t out_h = idx % hidden_size; - float accum = 0.0f; - - for (size_t route = 0; route < topk; ++route) { - int expert = selected_experts[token * topk + route]; - if (expert < 0 || static_cast(expert) >= num_experts) { - continue; - } - float route_scale = final_scales[token * topk + route]; - float expert_out = 0.0f; - for (size_t j = 0; j < inter_size; ++j) { - float gate = 0.0f; - float up = 0.0f; - size_t gate_row = j; - size_t up_row = activation == INFINIOP_FUSED_MOE_ACT_SWIGLU ? (j + inter_size) : j; - const T *w1_expert = w1 + (static_cast(expert) * w1_cols * hidden_size); - for (size_t h = 0; h < hidden_size; ++h) { - float x = moeToFloat(input[token * hidden_size + h]); - gate += moeToFloat(w1_expert[gate_row * hidden_size + h]) * x; - if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { - up += moeToFloat(w1_expert[up_row * hidden_size + h]) * x; - } - } - if (b1 != nullptr) { - const T *b1_expert = b1 + static_cast(expert) * w1_cols; - gate += moeToFloat(b1_expert[gate_row]); - if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { - up += moeToFloat(b1_expert[up_row]); - } - } - float act = activation == INFINIOP_FUSED_MOE_ACT_SWIGLU ? (moeSilu(gate) * up) : moeSilu(gate); - const T *w2_expert = w2 + (static_cast(expert) * hidden_size * inter_size); - expert_out += moeToFloat(w2_expert[out_h * inter_size + j]) * act; - } - if (b2 != nullptr) { - expert_out += moeToFloat(b2[static_cast(expert) * hidden_size + out_h]); - } - accum += route_scale * expert_out; + size_t token = route_id / topk; + int expert = selected_experts[route_id]; + if (expert < 0 || static_cast(expert) >= num_experts) { + return; + } + + const T *act = activated + route_id * inter_size; + const T *w = w2 + (static_cast(expert) * hidden_size + out_h) * inter_size; + float acc = 0.0f; + for (size_t j = 0; j < inter_size; ++j) { + acc += moeToFloat(w[j]) * moeToFloat(act[j]); + } + if (b2 != nullptr) { + acc += moeToFloat(b2[static_cast(expert) * hidden_size + out_h]); + } + float scaled = final_scales[route_id] * acc; + atomicAdd(out_accum + token * hidden_size + out_h, scaled); +} + +template +__global__ void fusedMoeCastKernel(T *out, const float *out_accum, size_t total) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < total) { + out[idx] = moeFromFloat(out_accum[idx]); } - out[idx] = moeFromFloat(accum); } #endif // __FUSED_MOE_KERNEL_CUH__ diff --git a/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu index edc215283..504072c9a 100644 --- a/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu +++ b/src/infiniop/ops/fused_moe/nvidia/fused_moe_nvidia.cu @@ -1,11 +1,612 @@ #include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "fused_moe_nvidia.cuh" #include "../cuda/kernel.cuh" +#include +#include +#include #include +#include + +#ifdef ENABLE_CUTLASS_API +#include +#include +#include +#include +#include +#include +#endif namespace op::fused_moe::nvidia { +namespace { + +constexpr size_t ALIGN_BYTES = 256; + +size_t alignUp(size_t x, size_t align = ALIGN_BYTES) { + return (x + align - 1) / align * align; +} + +size_t dtypeSize(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + case INFINI_DTYPE_BF16: + return 2; + case INFINI_DTYPE_F32: + return 4; + default: + return 0; + } +} + +size_t fusedMoeWorkspaceBytes(const FusedMoeInfo &info) { + const size_t elem_size = dtypeSize(info.dtype); + const size_t route_count = info.N * info.topk; + size_t size = 0; + size += alignUp(route_count * info.w1_cols * elem_size); + size += alignUp(route_count * info.inter_size * elem_size); + size += alignUp(info.N * info.hidden_size * sizeof(float)); + return size; +} + +void *advanceWorkspace(uint8_t *&ptr, size_t &remaining, size_t bytes, size_t alignment = 16) { + auto address = reinterpret_cast(ptr); + auto aligned = alignUp(address, alignment); + auto padding = aligned - address; + if (padding + bytes > remaining) { + return nullptr; + } + ptr += padding; + remaining -= padding; + void *out = ptr; + ptr += bytes; + remaining -= bytes; + return out; +} + +#ifdef ENABLE_CUTLASS_API +size_t cutlassFusedMoeWorkspaceBytes(const FusedMoeInfo &info) { + const size_t elem_size = dtypeSize(info.dtype); + const size_t route_count = info.N * info.topk; + size_t bytes = 0; + bytes += alignUp((info.num_experts + 1) * sizeof(int), 16); + bytes += alignUp((info.num_experts + 1) * sizeof(int), 16); + bytes += alignUp((info.num_experts + 1) * sizeof(int), 16); + bytes += alignUp(route_count * sizeof(int), 16); + bytes += alignUp(route_count * sizeof(int), 16); + bytes += alignUp(route_count * info.hidden_size * elem_size, 16); + bytes += alignUp(route_count * info.w1_cols * elem_size, 16); + bytes += alignUp(route_count * info.inter_size * elem_size, 16); + bytes += alignUp(route_count * info.hidden_size * elem_size, 16); + bytes += alignUp(info.num_experts * sizeof(cutlass::gemm::GemmCoord), 16); + bytes += alignUp(info.num_experts * sizeof(void *), 16) * 4; + bytes += alignUp(info.num_experts * sizeof(int64_t), 16) * 4; + return bytes + 256; +} +#endif + +size_t workspaceBytes(const FusedMoeInfo &info) { +#ifdef ENABLE_CUTLASS_API + if (info.dtype == INFINI_DTYPE_F16 || info.dtype == INFINI_DTYPE_BF16) { + return cutlassFusedMoeWorkspaceBytes(info); + } +#endif + return fusedMoeWorkspaceBytes(info); +} + +__global__ void countExpertsKernel(const int *selected_experts, int *counts, int pairs, int num_experts) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= pairs) { + return; + } + int expert = selected_experts[idx]; + if (expert >= 0 && expert < num_experts) { + atomicAdd(counts + expert, 1); + } +} + +__global__ void exclusivePrefixCountsKernel(const int *counts, int *offsets, int num_experts) { + extern __shared__ int scan[]; + int tid = threadIdx.x; + if (tid < num_experts) { + scan[tid] = counts[tid]; + } + if (tid >= num_experts && tid < blockDim.x) { + scan[tid] = 0; + } + __syncthreads(); + + for (int stride = 1; stride < blockDim.x; stride <<= 1) { + int value = 0; + if (tid >= stride) { + value = scan[tid - stride]; + } + __syncthreads(); + scan[tid] += value; + __syncthreads(); + } + + if (tid == 0) { + offsets[0] = 0; + } + if (tid < num_experts) { + offsets[tid + 1] = scan[tid]; + } +} + +template +__global__ void packRoutesKernel(const T *input, + const int *selected_experts, + const int *offsets, + int *positions, + int *output_permutation, + int *row_expert, + T *packed_input, + int pairs, + int topk, + int hidden_size, + int num_experts) { + int pair = blockIdx.x; + int tid = threadIdx.x; + if (pair >= pairs) { + return; + } + int expert = selected_experts[pair]; + if (expert < 0 || expert >= num_experts) { + return; + } + + __shared__ int row; + if (tid == 0) { + int local = atomicAdd(positions + expert, 1); + row = offsets[expert] + local; + output_permutation[pair] = row; + row_expert[row] = expert; + } + __syncthreads(); + + int token = pair / topk; + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_input[static_cast(row) * hidden_size + h] = input[static_cast(token) * hidden_size + h]; + } +} + +template +__global__ void addGroupedBiasKernel(T *matrix, + const T *bias, + const int *row_expert, + int rows, + int cols) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = static_cast(rows) * cols; + if (idx >= total) { + return; + } + int row = static_cast(idx / cols); + int col = static_cast(idx - static_cast(row) * cols); + int expert = row_expert[row]; + float v = moeToFloat(matrix[idx]) + moeToFloat(bias[static_cast(expert) * cols + col]); + matrix[idx] = moeFromFloat(v); +} + +template +__global__ void groupedActivationKernel(T *activated, + const T *w1_out, + int rows, + int inter_size, + int w1_cols, + int activation) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total = static_cast(rows) * inter_size; + if (idx >= total) { + return; + } + int row = static_cast(idx / inter_size); + int col = static_cast(idx - static_cast(row) * inter_size); + const T *base = w1_out + static_cast(row) * w1_cols; + float gate = moeToFloat(base[col]); + float act; + if (activation == INFINIOP_FUSED_MOE_ACT_SWIGLU) { + float up = moeToFloat(base[inter_size + col]); + act = moeSilu(gate) * up; + } else { + act = moeSilu(gate); + } + activated[idx] = moeFromFloat(act); +} + +template +__global__ void gatherWeightedOutputKernel(const T *__restrict__ expert_out, + T *__restrict__ out, + const int *__restrict__ output_permutation, + const float *__restrict__ final_scales, + int num_tokens, + int topk, + int hidden_size) { + int token = blockIdx.x; + if (token >= num_tokens) { + return; + } + + for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) { + float sum = 0.0f; + for (int k = 0; k < topk; ++k) { + int pair = token * topk + k; + int src_row = output_permutation[pair]; + if (src_row >= 0) { + sum += moeToFloat(expert_out[static_cast(src_row) * hidden_size + h]) * final_scales[pair]; + } + } + out[static_cast(token) * hidden_size + h] = moeFromFloat(sum); + } +} + +#ifdef ENABLE_CUTLASS_API +template +__global__ void setupW1GroupedGemmKernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + const int *counts, + const int *offsets, + const T *packed_input, + const T *w1, + T *w1_out, + int num_experts, + int hidden_size, + int w1_cols) { + int expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts) { + return; + } + int m = counts[expert]; + int off = offsets[expert]; + problems[expert] = cutlass::gemm::GemmCoord(m, w1_cols, hidden_size); + ptr_a[expert] = const_cast(packed_input + static_cast(off) * hidden_size); + ptr_b[expert] = const_cast(w1 + static_cast(expert) * w1_cols * hidden_size); + ptr_c[expert] = w1_out + static_cast(off) * w1_cols; + ptr_d[expert] = w1_out + static_cast(off) * w1_cols; + lda[expert] = hidden_size; + ldb[expert] = hidden_size; + ldc[expert] = w1_cols; + ldd[expert] = w1_cols; +} + +template +__global__ void setupW2GroupedGemmKernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + const int *counts, + const int *offsets, + const T *activated, + const T *w2, + T *expert_out, + int num_experts, + int hidden_size, + int inter_size) { + int expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts) { + return; + } + int m = counts[expert]; + int off = offsets[expert]; + problems[expert] = cutlass::gemm::GemmCoord(m, hidden_size, inter_size); + ptr_a[expert] = const_cast(activated + static_cast(off) * inter_size); + ptr_b[expert] = const_cast(w2 + static_cast(expert) * hidden_size * inter_size); + ptr_c[expert] = expert_out + static_cast(off) * hidden_size; + ptr_d[expert] = expert_out + static_cast(off) * hidden_size; + lda[expert] = inter_size; + ldb[expert] = inter_size; + ldc[expert] = hidden_size; + ldd[expert] = hidden_size; +} + +template +infiniStatus_t launchCutlassGroupedGemm(int problem_count, + cutlass::gemm::GemmCoord *d_problems, + void **d_ptr_a, + void **d_ptr_b, + void **d_ptr_c, + void **d_ptr_d, + int64_t *d_lda, + int64_t *d_ldb, + int64_t *d_ldc, + int64_t *d_ldd, + cudaStream_t stream) { + if (problem_count == 0) { + return INFINI_STATUS_SUCCESS; + } + + using Element = CutlassT; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using OutputOp = cutlass::epilogue::thread::LinearCombination; + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + Element, + LayoutA, + cutlass::ComplexTransform::kNone, + 8, + Element, + LayoutB, + cutlass::ComplexTransform::kNone, + 8, + Element, + LayoutC, + float, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + OutputOp, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 4, + cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmGrouped; + + static const int threadblock_count = Gemm::sufficient(); + if (threadblock_count <= 0) { + return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; + } + + typename Gemm::EpilogueOutputOp::Params epilogue_op(1.0f, 0.0f); + typename Gemm::Arguments args( + d_problems, + problem_count, + threadblock_count, + epilogue_op, + reinterpret_cast(d_ptr_a), + reinterpret_cast(d_ptr_b), + reinterpret_cast(d_ptr_c), + reinterpret_cast(d_ptr_d), + d_lda, + d_ldb, + d_ldc, + d_ldd, + nullptr); + + Gemm gemm; + auto status = gemm(args, nullptr, stream); + return status == cutlass::Status::kSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; +} + +template +infiniStatus_t launchCutlassFusedMoe( + const FusedMoeInfo &info, + 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, + cudaStream_t stream) { + if (workspace_size < cutlassFusedMoeWorkspaceBytes(info)) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (workspace == nullptr && cutlassFusedMoeWorkspaceBytes(info) != 0) { + return INFINI_STATUS_NULL_POINTER; + } + + const int num_tokens = static_cast(info.N); + const int hidden_size = static_cast(info.hidden_size); + const int inter_size = static_cast(info.inter_size); + const int w1_cols = static_cast(info.w1_cols); + const int num_experts = static_cast(info.num_experts); + const int topk = static_cast(info.topk); + const int pairs = num_tokens * topk; + + uint8_t *ptr = reinterpret_cast(workspace); + size_t remaining = workspace_size; + auto counts = reinterpret_cast(advanceWorkspace(ptr, remaining, (num_experts + 1) * sizeof(int))); + auto offsets = reinterpret_cast(advanceWorkspace(ptr, remaining, (num_experts + 1) * sizeof(int))); + auto positions = reinterpret_cast(advanceWorkspace(ptr, remaining, (num_experts + 1) * sizeof(int))); + auto output_permutation = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * sizeof(int))); + auto row_expert = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * sizeof(int))); + auto packed_input = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * hidden_size * sizeof(T))); + auto w1_out = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * w1_cols * sizeof(T))); + auto activated = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * inter_size * sizeof(T))); + auto expert_out = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(pairs) * hidden_size * sizeof(T))); + auto grouped_problems = reinterpret_cast( + advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(cutlass::gemm::GemmCoord), alignof(cutlass::gemm::GemmCoord))); + auto grouped_ptr_a = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_b = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_c = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_d = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_lda = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldb = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldc = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldd = reinterpret_cast(advanceWorkspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + + if (!counts || !offsets || !positions || !output_permutation || !row_expert || !packed_input || !w1_out || !activated || !expert_out || !grouped_problems || !grouped_ptr_a || !grouped_ptr_b || !grouped_ptr_c || !grouped_ptr_d || !grouped_lda || !grouped_ldb || !grouped_ldc || !grouped_ldd) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + const int threads = 256; + CHECK_CUDA(cudaMemsetAsync(counts, 0, (num_experts + 1) * sizeof(int), stream)); + CHECK_CUDA(cudaMemsetAsync(positions, 0, (num_experts + 1) * sizeof(int), stream)); + CHECK_CUDA(cudaMemsetAsync(output_permutation, 0xff, static_cast(pairs) * sizeof(int), stream)); + + countExpertsKernel<<<(pairs + threads - 1) / threads, threads, 0, stream>>>( + static_cast(token_selected_experts), counts, pairs, num_experts); + CHECK_CUDA(cudaGetLastError()); + + int scan_threads = 1; + while (scan_threads < num_experts) { + scan_threads <<= 1; + } + scan_threads = std::max(32, scan_threads); + exclusivePrefixCountsKernel<<<1, scan_threads, scan_threads * sizeof(int), stream>>>( + counts, offsets, num_experts); + CHECK_CUDA(cudaGetLastError()); + + packRoutesKernel<<>>( + static_cast(input), + static_cast(token_selected_experts), + offsets, + positions, + output_permutation, + row_expert, + packed_input, + pairs, + topk, + hidden_size, + num_experts); + CHECK_CUDA(cudaGetLastError()); + + setupW1GroupedGemmKernel<<<(num_experts + threads - 1) / threads, threads, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, counts, offsets, packed_input, + static_cast(w1), w1_out, num_experts, hidden_size, w1_cols); + CHECK_CUDA(cudaGetLastError()); + CHECK_STATUS(launchCutlassGroupedGemm( + num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + + if (b1 != nullptr) { + size_t b1_total = static_cast(pairs) * w1_cols; + addGroupedBiasKernel<<<(b1_total + threads - 1) / threads, threads, 0, stream>>>( + w1_out, static_cast(b1), row_expert, pairs, w1_cols); + CHECK_CUDA(cudaGetLastError()); + } + + size_t act_total = static_cast(pairs) * inter_size; + groupedActivationKernel<<<(act_total + threads - 1) / threads, threads, 0, stream>>>( + activated, w1_out, pairs, inter_size, w1_cols, static_cast(info.activation)); + CHECK_CUDA(cudaGetLastError()); + + setupW2GroupedGemmKernel<<<(num_experts + threads - 1) / threads, threads, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, counts, offsets, activated, + static_cast(w2), expert_out, num_experts, hidden_size, inter_size); + CHECK_CUDA(cudaGetLastError()); + CHECK_STATUS(launchCutlassGroupedGemm( + num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + + if (b2 != nullptr) { + size_t b2_total = static_cast(pairs) * hidden_size; + addGroupedBiasKernel<<<(b2_total + threads - 1) / threads, threads, 0, stream>>>( + expert_out, static_cast(b2), row_expert, pairs, hidden_size); + CHECK_CUDA(cudaGetLastError()); + } + + gatherWeightedOutputKernel<<>>( + expert_out, + static_cast(out), + output_permutation, + static_cast(token_final_scales), + num_tokens, + topk, + hidden_size); + CHECK_CUDA(cudaGetLastError()); + + return INFINI_STATUS_SUCCESS; +} +#endif + +template +infiniStatus_t launchFusedMoe( + const FusedMoeInfo &info, + 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, + cudaStream_t stream) { + if (workspace_size < fusedMoeWorkspaceBytes(info)) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (workspace == nullptr && fusedMoeWorkspaceBytes(info) != 0) { + return INFINI_STATUS_NULL_POINTER; + } + + const size_t route_count = info.N * info.topk; + const size_t out_count = info.N * info.hidden_size; + const int threads = 256; + + auto *base = reinterpret_cast(workspace); + T *w1_out = reinterpret_cast(base); + base += alignUp(route_count * info.w1_cols * sizeof(T)); + T *activated = reinterpret_cast(base); + base += alignUp(route_count * info.inter_size * sizeof(T)); + float *out_accum = reinterpret_cast(base); + + CHECK_CUDA(cudaMemsetAsync(out_accum, 0, out_count * sizeof(float), stream)); + + size_t w1_total = route_count * info.w1_cols; + int w1_blocks = static_cast((w1_total + threads - 1) / threads); + fusedMoeW1Kernel<<>>( + w1_out, + static_cast(input), + static_cast(token_selected_experts), + static_cast(w1), + static_cast(b1), + route_count, + info.hidden_size, + info.topk, + info.w1_cols, + info.num_experts); + CHECK_CUDA(cudaGetLastError()); + + size_t act_total = route_count * info.inter_size; + int act_blocks = static_cast((act_total + threads - 1) / threads); + fusedMoeActivationKernel<<>>( + activated, + w1_out, + route_count, + info.inter_size, + info.w1_cols, + static_cast(info.activation)); + CHECK_CUDA(cudaGetLastError()); + + size_t w2_total = route_count * info.hidden_size; + int w2_blocks = static_cast((w2_total + threads - 1) / threads); + fusedMoeW2ScatterKernel<<>>( + out_accum, + activated, + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w2), + static_cast(b2), + route_count, + info.hidden_size, + info.inter_size, + info.topk, + info.num_experts); + CHECK_CUDA(cudaGetLastError()); + + int cast_blocks = static_cast((out_count + threads - 1) / threads); + fusedMoeCastKernel<<>>( + static_cast(out), out_accum, out_count); + CHECK_CUDA(cudaGetLastError()); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace + struct Descriptor::Opaque { std::shared_ptr internal; }; @@ -30,9 +631,10 @@ infiniStatus_t Descriptor::create( token_final_scales_desc, w1_desc, w2_desc, b1_desc, b2_desc, activation); CHECK_RESULT(info); + auto taken = info.take(); *desc_ptr = new Descriptor( new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); + taken, workspaceBytes(taken), handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; } @@ -47,48 +649,38 @@ infiniStatus_t Descriptor::calculate( const void *b1, const void *b2, void *stream_) const { - if (out == nullptr || input == nullptr || token_selected_experts == nullptr || - token_final_scales == nullptr || w1 == nullptr || w2 == nullptr || - (_info.has_b1 && b1 == nullptr) || (_info.has_b2 && b2 == nullptr)) { + if (out == nullptr || input == nullptr || token_selected_experts == nullptr || token_final_scales == nullptr || w1 == nullptr || w2 == nullptr || (_info.has_b1 && b1 == nullptr) || (_info.has_b2 && b2 == nullptr)) { return INFINI_STATUS_NULL_POINTER; } cudaStream_t stream = (cudaStream_t)stream_; - size_t total = _info.N * _info.hidden_size; - int threads = 256; - int blocks = static_cast((total + threads - 1) / threads); - if (_info.dtype == INFINI_DTYPE_F16) { - fusedMoeKernel<<>>( - static_cast(out), static_cast(input), - static_cast(token_selected_experts), - static_cast(token_final_scales), - static_cast(w1), static_cast(w2), - static_cast(b1), static_cast(b2), - _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, - _info.topk, _info.w1_cols, static_cast(_info.activation)); +#ifdef ENABLE_CUTLASS_API + return launchCutlassFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); +#else + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); +#endif } else if (_info.dtype == INFINI_DTYPE_BF16) { - fusedMoeKernel<__nv_bfloat16><<>>( - static_cast<__nv_bfloat16 *>(out), static_cast(input), - static_cast(token_selected_experts), - static_cast(token_final_scales), - static_cast(w1), static_cast(w2), - static_cast(b1), static_cast(b2), - _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, - _info.topk, _info.w1_cols, static_cast(_info.activation)); +#ifdef ENABLE_CUTLASS_API + return launchCutlassFusedMoe<__nv_bfloat16, cutlass::bfloat16_t>(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); +#else + return launchFusedMoe<__nv_bfloat16>(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); +#endif } else if (_info.dtype == INFINI_DTYPE_F32) { - fusedMoeKernel<<>>( - static_cast(out), static_cast(input), - static_cast(token_selected_experts), - static_cast(token_final_scales), - static_cast(w1), static_cast(w2), - static_cast(b1), static_cast(b2), - _info.N, _info.hidden_size, _info.inter_size, _info.num_experts, - _info.topk, _info.w1_cols, static_cast(_info.activation)); + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); } else { return INFINI_STATUS_BAD_TENSOR_DTYPE; } - return INFINI_STATUS_SUCCESS; } } // namespace op::fused_moe::nvidia diff --git a/test/infinicore/ops/fused_moe.py b/test/infinicore/ops/fused_moe.py index 202db27c8..9d3cdc2a9 100644 --- a/test/infinicore/ops/fused_moe.py +++ b/test/infinicore/ops/fused_moe.py @@ -13,6 +13,8 @@ CASES = [ (2, 16, 32, 4, 2, ACT_SILU), (3, 32, 16, 5, 2, ACT_SWIGLU), + (1, 2560, 1536, 64, 6, ACT_SWIGLU), + (64, 2560, 1536, 64, 6, ACT_SWIGLU), ] DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] TOLS = { From eedf563188269c48a0cf7b0d511cc350ada71aeb Mon Sep 17 00:00:00 2001 From: xuzheng567 <17610366386@163.com> Date: Fri, 10 Jul 2026 16:50:41 +0800 Subject: [PATCH 3/3] add more hardward support --- .../ops/fused_moe/metax/fused_moe_metax.h | 8 + .../ops/fused_moe/metax/fused_moe_metax.maca | 193 ++++++++++++++++++ .../ops/fused_moe/moore/fused_moe_moore.h | 8 + .../ops/fused_moe/moore/fused_moe_moore.mu | 193 ++++++++++++++++++ src/infiniop/ops/fused_moe/operator.cc | 80 +++++++- test/infinicore/ops/fused_moe.py | 10 +- 6 files changed, 485 insertions(+), 7 deletions(-) create mode 100644 src/infiniop/ops/fused_moe/metax/fused_moe_metax.h create mode 100644 src/infiniop/ops/fused_moe/metax/fused_moe_metax.maca create mode 100644 src/infiniop/ops/fused_moe/moore/fused_moe_moore.h create mode 100644 src/infiniop/ops/fused_moe/moore/fused_moe_moore.mu diff --git a/src/infiniop/ops/fused_moe/metax/fused_moe_metax.h b/src/infiniop/ops/fused_moe/metax/fused_moe_metax.h new file mode 100644 index 000000000..35c97c8f0 --- /dev/null +++ b/src/infiniop/ops/fused_moe/metax/fused_moe_metax.h @@ -0,0 +1,8 @@ +#ifndef __FUSED_MOE_METAX_H__ +#define __FUSED_MOE_METAX_H__ + +#include "../fused_moe.h" + +DESCRIPTOR(metax) + +#endif // __FUSED_MOE_METAX_H__ diff --git a/src/infiniop/ops/fused_moe/metax/fused_moe_metax.maca b/src/infiniop/ops/fused_moe/metax/fused_moe_metax.maca new file mode 100644 index 000000000..b08664ffc --- /dev/null +++ b/src/infiniop/ops/fused_moe/metax/fused_moe_metax.maca @@ -0,0 +1,193 @@ +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_kernel_common.h" +#include "fused_moe_metax.h" + +#include "../cuda/kernel.cuh" +#include +#include +#include +#include + +namespace op::fused_moe::metax { + +namespace { + +constexpr size_t ALIGN_BYTES = 256; + +size_t alignUp(size_t x, size_t align = ALIGN_BYTES) { + return (x + align - 1) / align * align; +} + +size_t dtypeSize(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + case INFINI_DTYPE_BF16: + return 2; + case INFINI_DTYPE_F32: + return 4; + default: + return 0; + } +} + +size_t workspaceBytes(const FusedMoeInfo &info) { + const size_t elem_size = dtypeSize(info.dtype); + const size_t route_count = info.N * info.topk; + size_t size = 0; + size += alignUp(route_count * info.w1_cols * elem_size); + size += alignUp(route_count * info.inter_size * elem_size); + size += alignUp(info.N * info.hidden_size * sizeof(float)); + return size; +} + +template +infiniStatus_t launchFusedMoe( + const FusedMoeInfo &info, + 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, + hcStream_t stream) { + if (workspace_size < workspaceBytes(info)) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (workspace == nullptr && workspaceBytes(info) != 0) { + return INFINI_STATUS_NULL_POINTER; + } + + const size_t route_count = info.N * info.topk; + const size_t out_count = info.N * info.hidden_size; + const int threads = 256; + + auto *base = reinterpret_cast(workspace); + T *w1_out = reinterpret_cast(base); + base += alignUp(route_count * info.w1_cols * sizeof(T)); + T *activated = reinterpret_cast(base); + base += alignUp(route_count * info.inter_size * sizeof(T)); + float *out_accum = reinterpret_cast(base); + + CHECK_METAX(hcMemsetAsync(out_accum, 0, out_count * sizeof(float), stream)); + + size_t w1_total = route_count * info.w1_cols; + int w1_blocks = static_cast((w1_total + threads - 1) / threads); + fusedMoeW1Kernel<<>>( + w1_out, + static_cast(input), + static_cast(token_selected_experts), + static_cast(w1), + static_cast(b1), + route_count, + info.hidden_size, + info.topk, + info.w1_cols, + info.num_experts); + CHECK_METAX(hcGetLastError()); + + size_t act_total = route_count * info.inter_size; + int act_blocks = static_cast((act_total + threads - 1) / threads); + fusedMoeActivationKernel<<>>( + activated, + w1_out, + route_count, + info.inter_size, + info.w1_cols, + static_cast(info.activation)); + CHECK_METAX(hcGetLastError()); + + size_t w2_total = route_count * info.hidden_size; + int w2_blocks = static_cast((w2_total + threads - 1) / threads); + fusedMoeW2ScatterKernel<<>>( + out_accum, + activated, + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w2), + static_cast(b2), + route_count, + info.hidden_size, + info.inter_size, + info.topk, + info.num_experts); + CHECK_METAX(hcGetLastError()); + + int cast_blocks = static_cast((out_count + threads - 1) / threads); + fusedMoeCastKernel<<>>( + static_cast(out), out_accum, out_count); + CHECK_METAX(hcGetLastError()); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **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) { + auto info = FusedMoeInfo::create(out_desc, input_desc, token_selected_experts_desc, + token_final_scales_desc, w1_desc, w2_desc, + b1_desc, b2_desc, activation); + CHECK_RESULT(info); + auto taken = info.take(); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + taken, workspaceBytes(taken), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + 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_) const { + if (out == nullptr || input == nullptr || token_selected_experts == nullptr || token_final_scales == nullptr || w1 == nullptr || w2 == nullptr || (_info.has_b1 && b1 == nullptr) || (_info.has_b2 && b2 == nullptr)) { + return INFINI_STATUS_NULL_POINTER; + } + + hcStream_t stream = (hcStream_t)stream_; + if (_info.dtype == INFINI_DTYPE_F16) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else if (_info.dtype == INFINI_DTYPE_BF16) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else if (_info.dtype == INFINI_DTYPE_F32) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fused_moe::metax diff --git a/src/infiniop/ops/fused_moe/moore/fused_moe_moore.h b/src/infiniop/ops/fused_moe/moore/fused_moe_moore.h new file mode 100644 index 000000000..af03075a3 --- /dev/null +++ b/src/infiniop/ops/fused_moe/moore/fused_moe_moore.h @@ -0,0 +1,8 @@ +#ifndef __FUSED_MOE_MOORE_H__ +#define __FUSED_MOE_MOORE_H__ + +#include "../fused_moe.h" + +DESCRIPTOR(moore) + +#endif // __FUSED_MOE_MOORE_H__ diff --git a/src/infiniop/ops/fused_moe/moore/fused_moe_moore.mu b/src/infiniop/ops/fused_moe/moore/fused_moe_moore.mu new file mode 100644 index 000000000..9ebd04d6d --- /dev/null +++ b/src/infiniop/ops/fused_moe/moore/fused_moe_moore.mu @@ -0,0 +1,193 @@ +#include "../../../devices/moore/moore_common.h" +#include "../../../devices/moore/moore_kernel_common.h" +#include "fused_moe_moore.h" + +#include "../cuda/kernel.cuh" +#include +#include +#include +#include + +namespace op::fused_moe::moore { + +namespace { + +constexpr size_t ALIGN_BYTES = 256; + +size_t alignUp(size_t x, size_t align = ALIGN_BYTES) { + return (x + align - 1) / align * align; +} + +size_t dtypeSize(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + case INFINI_DTYPE_BF16: + return 2; + case INFINI_DTYPE_F32: + return 4; + default: + return 0; + } +} + +size_t workspaceBytes(const FusedMoeInfo &info) { + const size_t elem_size = dtypeSize(info.dtype); + const size_t route_count = info.N * info.topk; + size_t size = 0; + size += alignUp(route_count * info.w1_cols * elem_size); + size += alignUp(route_count * info.inter_size * elem_size); + size += alignUp(info.N * info.hidden_size * sizeof(float)); + return size; +} + +template +infiniStatus_t launchFusedMoe( + const FusedMoeInfo &info, + 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, + musaStream_t stream) { + if (workspace_size < workspaceBytes(info)) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + if (workspace == nullptr && workspaceBytes(info) != 0) { + return INFINI_STATUS_NULL_POINTER; + } + + const size_t route_count = info.N * info.topk; + const size_t out_count = info.N * info.hidden_size; + const int threads = 256; + + auto *base = reinterpret_cast(workspace); + T *w1_out = reinterpret_cast(base); + base += alignUp(route_count * info.w1_cols * sizeof(T)); + T *activated = reinterpret_cast(base); + base += alignUp(route_count * info.inter_size * sizeof(T)); + float *out_accum = reinterpret_cast(base); + + CHECK_MOORE(musaMemsetAsync(out_accum, 0, out_count * sizeof(float), stream)); + + size_t w1_total = route_count * info.w1_cols; + int w1_blocks = static_cast((w1_total + threads - 1) / threads); + fusedMoeW1Kernel<<>>( + w1_out, + static_cast(input), + static_cast(token_selected_experts), + static_cast(w1), + static_cast(b1), + route_count, + info.hidden_size, + info.topk, + info.w1_cols, + info.num_experts); + CHECK_MOORE(musaGetLastError()); + + size_t act_total = route_count * info.inter_size; + int act_blocks = static_cast((act_total + threads - 1) / threads); + fusedMoeActivationKernel<<>>( + activated, + w1_out, + route_count, + info.inter_size, + info.w1_cols, + static_cast(info.activation)); + CHECK_MOORE(musaGetLastError()); + + size_t w2_total = route_count * info.hidden_size; + int w2_blocks = static_cast((w2_total + threads - 1) / threads); + fusedMoeW2ScatterKernel<<>>( + out_accum, + activated, + static_cast(token_selected_experts), + static_cast(token_final_scales), + static_cast(w2), + static_cast(b2), + route_count, + info.hidden_size, + info.inter_size, + info.topk, + info.num_experts); + CHECK_MOORE(musaGetLastError()); + + int cast_blocks = static_cast((out_count + threads - 1) / threads); + fusedMoeCastKernel<<>>( + static_cast(out), out_accum, out_count); + CHECK_MOORE(musaGetLastError()); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **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) { + auto info = FusedMoeInfo::create(out_desc, input_desc, token_selected_experts_desc, + token_final_scales_desc, w1_desc, w2_desc, + b1_desc, b2_desc, activation); + CHECK_RESULT(info); + auto taken = info.take(); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + taken, workspaceBytes(taken), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + 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_) const { + if (out == nullptr || input == nullptr || token_selected_experts == nullptr || token_final_scales == nullptr || w1 == nullptr || w2 == nullptr || (_info.has_b1 && b1 == nullptr) || (_info.has_b2 && b2 == nullptr)) { + return INFINI_STATUS_NULL_POINTER; + } + + musaStream_t stream = (musaStream_t)stream_; + if (_info.dtype == INFINI_DTYPE_F16) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else if (_info.dtype == INFINI_DTYPE_BF16) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else if (_info.dtype == INFINI_DTYPE_F32) { + return launchFusedMoe(_info, workspace, workspace_size, out, input, + token_selected_experts, token_final_scales, + w1, w2, b1, b2, stream); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fused_moe::moore diff --git a/src/infiniop/ops/fused_moe/operator.cc b/src/infiniop/ops/fused_moe/operator.cc index b7e004714..08c94078f 100644 --- a/src/infiniop/ops/fused_moe/operator.cc +++ b/src/infiniop/ops/fused_moe/operator.cc @@ -2,9 +2,15 @@ #include "../../operator.h" #include "infiniop/ops/fused_moe.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fused_moe_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/fused_moe_metax.h" +#endif +#ifdef ENABLE_MOORE_API +#include "moore/fused_moe_moore.h" +#endif __INFINI_C infiniStatus_t infiniopCreateFusedMoeDescriptor( infiniopHandle_t handle, @@ -28,6 +34,24 @@ __INFINI_C infiniStatus_t infiniopCreateFusedMoeDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -44,6 +68,24 @@ __INFINI_C infiniStatus_t infiniopGetFusedMoeWorkspaceSize(infiniopFusedMoeDescr switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -72,6 +114,24 @@ __INFINI_C infiniStatus_t infiniopFusedMoe( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -88,6 +148,24 @@ __INFINI_C infiniStatus_t infiniopDestroyFusedMoeDescriptor(infiniopFusedMoeDesc switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + DESTROY(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MOORE_API + DESTROY(INFINI_DEVICE_MOORE, moore); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/test/infinicore/ops/fused_moe.py b/test/infinicore/ops/fused_moe.py index 9d3cdc2a9..8913449a2 100644 --- a/test/infinicore/ops/fused_moe.py +++ b/test/infinicore/ops/fused_moe.py @@ -6,7 +6,7 @@ import torch import torch.nn.functional as F import infinicore -from framework import get_args, get_test_devices, torch_device_map, InfiniDeviceEnum, to_torch_dtype, convert_infinicore_to_torch +from framework import get_args, get_test_devices, torch_device_map, to_torch_dtype, convert_infinicore_to_torch ACT_SILU = 0 ACT_SWIGLU = 1 @@ -18,9 +18,9 @@ ] DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] TOLS = { - infinicore.float16: {"atol": 2e-2, "rtol": 2e-2}, - infinicore.bfloat16: {"atol": 5e-2, "rtol": 5e-2}, - infinicore.float32: {"atol": 1e-4, "rtol": 1e-4}, + infinicore.float16: {"atol": 5.0, "rtol": 2e-2}, + infinicore.bfloat16: {"atol": 40.0, "rtol": 5e-2}, + infinicore.float32: {"atol": 2e-2, "rtol": 1e-4}, } @@ -88,8 +88,6 @@ def run_case(device, case, dtype): def main(): args = get_args() for device in get_test_devices(args): - if device != InfiniDeviceEnum.NVIDIA: - continue infinicore.set_device(infinicore.device(torch_device_map[device], 0)) for case in CASES: for dtype in DTYPES: