feat(nvidia): add moe_sum operator

src/CMakeLists.txt

@@ -472,7 +472,7 @@ set(_infini_ops_smoke_ops
 set(_infini_ops_smoke_torch_ops abs clamp exp)
 
 if(WITH_NVIDIA)
-    list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm)
+    list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum)
 endif()
 
 if(WITH_NVIDIA AND WITH_TORCH)

src/base/moe_sum.h

@@ -0,0 +1,213 @@
+#ifndef INFINI_OPS_BASE_MOE_SUM_H_
+#define INFINI_OPS_BASE_MOE_SUM_H_
+
+#include <cassert>
+#include <cstdint>
+#include <limits>
+#include <optional>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Aligned with vLLM `_moe_C::moe_sum`.
+class MoeSum : public Operator<MoeSum> {
+ public:
+  MoeSum(const Tensor input, Tensor output)
+      : MoeSum{input, std::nullopt, std::nullopt, output} {}
+
+  MoeSum(const Tensor input, std::optional<Tensor> topk_ids,
+         std::optional<Tensor> expert_map, Tensor output)
+      : num_tokens_{input.ndim() == 3 ? input.size(0) : 0},
+        topk_{input.ndim() == 3 ? input.size(1) : 0},
+        hidden_size_{input.ndim() == 3 ? input.size(2) : 0},
+        input_strides_{input.strides()},
+        output_strides_{output.strides()},
+        dtype_{input.dtype()},
+        device_type_{input.device().type()},
+        has_topk_ids_{topk_ids.has_value()},
+        topk_ids_dtype_{topk_ids ? topk_ids->dtype() : DataType::kInt32},
+        topk_ids_token_stride_{
+            topk_ids && topk_ids->ndim() == 2 ? topk_ids->stride(0) : 0},
+        topk_ids_slot_stride_{
+            topk_ids && topk_ids->ndim() == 2 ? topk_ids->stride(1) : 0},
+        has_expert_map_{expert_map.has_value()},
+        expert_map_size_{expert_map ? expert_map->numel() : 0},
+        expert_map_stride_{
+            expert_map && expert_map->ndim() == 1 ? expert_map->stride(0) : 0},
+        device_index_{input.device().index()} {
+    assert(input.ndim() == 3 && output.ndim() == 2 &&
+           "`MoeSum` requires `[num_tokens, topk, hidden_size]` input and "
+           "`[num_tokens, hidden_size]` output");
+    assert(output.size(0) == num_tokens_ && output.size(1) == hidden_size_ &&
+           "`MoeSum` output shape is incompatible with the input");
+    assert(topk_ > 0 && "`MoeSum` requires at least one top-k slot");
+    assert((dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat16 ||
+            dtype_ == DataType::kBFloat16) &&
+           "`MoeSum` supports float32, float16, and bfloat16 inputs");
+    assert(output.dtype() == dtype_ &&
+           "`MoeSum` input and output dtypes must match");
+    assert(output.IsContiguous() && "`MoeSum` requires contiguous output");
+
+    constexpr auto kMaxSignedIndex =
+        static_cast<Tensor::Size>(std::numeric_limits<int64_t>::max());
+    assert(num_tokens_ <= kMaxSignedIndex && topk_ <= kMaxSignedIndex &&
+           hidden_size_ <= kMaxSignedIndex &&
+           "`MoeSum` dimensions must fit signed index arithmetic");
+    assert(num_tokens_ <= std::numeric_limits<int>::max() &&
+           "`MoeSum` token count exceeds the CUDA grid limit");
+    assert(
+        (hidden_size_ == 0 || num_tokens_ <= kMaxSignedIndex / hidden_size_) &&
+        "`MoeSum` output size must fit signed index arithmetic");
+
+    const auto offsets_fit = [](const Tensor::Shape& shape,
+                                const Tensor::Strides& strides) {
+      uint64_t max_offset = 0;
+      constexpr auto kMaxOffset =
+          static_cast<uint64_t>(std::numeric_limits<int64_t>::max());
+
+      for (Tensor::Size dim = 0; dim < shape.size(); ++dim) {
+        if (static_cast<uint64_t>(shape[dim]) > kMaxOffset) {
+          return false;
+        }
+        if (strides[dim] < 0) {
+          return false;
+        }
+        if (shape[dim] == 0) {
+          continue;
+        }
+
+        const auto extent = static_cast<uint64_t>(shape[dim] - 1);
+        const auto stride = static_cast<uint64_t>(strides[dim]);
+        if (extent != 0 && stride > kMaxOffset / extent) {
+          return false;
+        }
+
+        const auto term = extent * stride;
+        if (term > kMaxOffset - max_offset) {
+          return false;
+        }
+        max_offset += term;
+      }
+
+      return true;
+    };
+    assert(offsets_fit(input.shape(), input.strides()) &&
+           "`MoeSum` input requires non-negative strides with signed offsets");
+
+    const auto same_device_as_input = [&](const Tensor tensor) {
+      return tensor.device().type() == input.device().type() &&
+             tensor.device().index() == input.device().index();
+    };
+    assert(same_device_as_input(output) &&
+           "`MoeSum` input and output must be on the same device");
+    assert((!expert_map || topk_ids) &&
+           "`MoeSum` expert_map requires topk_ids");
+
+    if (topk_ids) {
+      assert(topk_ids->ndim() == 2 && topk_ids->size(0) == num_tokens_ &&
+             topk_ids->size(1) == topk_ &&
+             "`MoeSum` topk_ids must have shape `[num_tokens, topk]`");
+      assert((topk_ids_dtype_ == DataType::kInt32 ||
+              topk_ids_dtype_ == DataType::kInt64) &&
+             "`MoeSum` topk_ids must have int32 or int64 dtype");
+      assert(same_device_as_input(*topk_ids) &&
+             "`MoeSum` topk_ids must be on the input device");
+      assert(offsets_fit(topk_ids->shape(), topk_ids->strides()) &&
+             "`MoeSum` topk_ids requires non-negative strides with signed "
+             "offsets");
+    }
+
+    if (expert_map) {
+      assert(expert_map->ndim() == 1 &&
+             expert_map->dtype() == DataType::kInt32 &&
+             "`MoeSum` expert_map must be a 1D int32 tensor");
+      assert(same_device_as_input(*expert_map) &&
+             "`MoeSum` expert_map must be on the input device");
+      assert(offsets_fit(expert_map->shape(), expert_map->strides()) &&
+             "`MoeSum` expert_map requires non-negative strides with signed "
+             "offsets");
+    }
+  }
+
+  void operator()(const Tensor input, Tensor output) const {
+    (*this)(input, std::nullopt, std::nullopt, output);
+  }
+
+  virtual void operator()(const Tensor input, std::optional<Tensor> topk_ids,
+                          std::optional<Tensor> expert_map,
+                          Tensor output) const = 0;
+
+ protected:
+  void ValidateCallMetadata(const Tensor input, std::optional<Tensor> topk_ids,
+                            std::optional<Tensor> expert_map,
+                            const Tensor output) const {
+    const auto same_device_as_descriptor = [&](const Tensor tensor) {
+      return tensor.device().type() == device_type_ &&
+             tensor.device().index() == device_index_;
+    };
+    auto matches =
+        input.ndim() == 3 && input.size(0) == num_tokens_ &&
+        input.size(1) == topk_ && input.size(2) == hidden_size_ &&
+        input.strides() == input_strides_ && input.dtype() == dtype_ &&
+        same_device_as_descriptor(input) && output.ndim() == 2 &&
+        output.size(0) == num_tokens_ && output.size(1) == hidden_size_ &&
+        output.strides() == output_strides_ && output.dtype() == dtype_ &&
+        same_device_as_descriptor(output) &&
+        topk_ids.has_value() == has_topk_ids_ &&
+        expert_map.has_value() == has_expert_map_;
+
+    if (matches && topk_ids) {
+      matches = topk_ids->ndim() == 2 && topk_ids->size(0) == num_tokens_ &&
+                topk_ids->size(1) == topk_ &&
+                topk_ids->stride(0) == topk_ids_token_stride_ &&
+                topk_ids->stride(1) == topk_ids_slot_stride_ &&
+                topk_ids->dtype() == topk_ids_dtype_ &&
+                same_device_as_descriptor(*topk_ids);
+    }
+
+    if (matches && expert_map) {
+      matches = expert_map->ndim() == 1 &&
+                expert_map->numel() == expert_map_size_ &&
+                expert_map->stride(0) == expert_map_stride_ &&
+                expert_map->dtype() == DataType::kInt32 &&
+                same_device_as_descriptor(*expert_map);
+    }
+
+    assert(matches && "`MoeSum` call metadata must match descriptor");
+  }
+
+  Tensor::Size num_tokens_{0};
+
+  Tensor::Size topk_{0};
+
+  Tensor::Size hidden_size_{0};
+
+  Tensor::Strides input_strides_;
+
+  Tensor::Strides output_strides_;
+
+  DataType dtype_;
+
+  Device::Type device_type_;
+
+  bool has_topk_ids_{false};
+
+  DataType topk_ids_dtype_;
+
+  Tensor::Stride topk_ids_token_stride_{0};
+
+  Tensor::Stride topk_ids_slot_stride_{0};
+
+  bool has_expert_map_{false};
+
+  Tensor::Size expert_map_size_{0};
+
+  Tensor::Stride expert_map_stride_{0};
+
+  int device_index_{0};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_BASE_MOE_SUM_H_

src/native/cuda/nvidia/ops/moe_sum/kernel.cu

@@ -0,0 +1,86 @@
+#include <cassert>
+#include <cstdint>
+
+#include "data_type.h"
+#include "dispatcher.h"
+#include "native/cuda/nvidia/caster.cuh"
+#include "native/cuda/nvidia/ops/moe_sum/kernel.h"
+#include "native/cuda/nvidia/runtime_.h"
+#include "native/cuda/ops/moe_sum/kernel.cuh"
+
+namespace infini::ops {
+namespace {
+
+class DeviceGuard {
+ public:
+  explicit DeviceGuard(int device_index) {
+    auto status = cudaGetDevice(&previous_device_);
+    assert(status == cudaSuccess &&
+           "`MoeSum` failed to query the current CUDA device");
+
+    if (previous_device_ != device_index) {
+      status = cudaSetDevice(device_index);
+      assert(status == cudaSuccess &&
+             "`MoeSum` failed to select the input CUDA device");
+      restore_ = true;
+    }
+  }
+
+  ~DeviceGuard() {
+    if (restore_) {
+      const auto status = cudaSetDevice(previous_device_);
+      assert(status == cudaSuccess &&
+             "`MoeSum` failed to restore the CUDA device");
+    }
+  }
+
+ private:
+  int previous_device_{0};
+
+  bool restore_{false};
+};
+
+}  // namespace
+
+void Operator<MoeSum, Device::Type::kNvidia, 0>::operator()(
+    const Tensor input, std::optional<Tensor> topk_ids,
+    std::optional<Tensor> expert_map, Tensor output) const {
+  ValidateCallMetadata(input, topk_ids, expert_map, output);
+
+  if (num_tokens_ == 0 || hidden_size_ == 0) {
+    return;
+  }
+
+  DeviceGuard device_guard{device_index_};
+  auto stream = static_cast<cudaStream_t>(stream_ ? stream_ : 0);
+  constexpr unsigned int kBlockSize = 256;
+  using DataTypes = ConcatType<List<DataType::kFloat32>, ReducedFloatTypes>;
+  using IndexTypes = List<DataType::kInt32, DataType::kInt64>;
+
+  DispatchFunc<DataTypes, IndexTypes>(
+      {static_cast<int64_t>(dtype_), static_cast<int64_t>(topk_ids_dtype_)},
+      [&](auto list_tag) {
+        using Data = TypeMapType<Device::Type::kNvidia, ListGet<0>(list_tag)>;
+        using Index = TypeMapType<Device::Type::kNvidia, ListGet<1>(list_tag)>;
+
+        MoeSumKernel<Device::Type::kNvidia, Data, Index>
+            <<<static_cast<unsigned int>(num_tokens_), kBlockSize, 0, stream>>>(
+                reinterpret_cast<const Data*>(input.data()),
+                topk_ids ? reinterpret_cast<const Index*>(topk_ids->data())
+                         : nullptr,
+                expert_map
+                    ? reinterpret_cast<const int32_t*>(expert_map->data())
+                    : nullptr,
+                reinterpret_cast<Data*>(output.data()),
+                static_cast<int64_t>(topk_), static_cast<int64_t>(hidden_size_),
+                input_strides_[0], input_strides_[1], input_strides_[2],
+                topk_ids_token_stride_, topk_ids_slot_stride_,
+                static_cast<int64_t>(expert_map_size_), expert_map_stride_);
+      },
+      "Operator<MoeSum, Device::Type::kNvidia>::operator()");
+
+  const auto status = cudaGetLastError();
+  assert(status == cudaSuccess && "`MoeSum` CUDA kernel launch failed");
+}
+
+}  // namespace infini::ops

src/native/cuda/nvidia/ops/moe_sum/kernel.h

@@ -0,0 +1,24 @@
+#ifndef INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_
+#define INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_
+
+#include <optional>
+
+#include "base/moe_sum.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<MoeSum, Device::Type::kNvidia, 0> : public MoeSum {
+ public:
+  using MoeSum::MoeSum;
+
+  using MoeSum::operator();
+
+  void operator()(const Tensor input, std::optional<Tensor> topk_ids,
+                  std::optional<Tensor> expert_map,
+                  Tensor output) const override;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_

src/native/cuda/ops/moe_sum/kernel.cuh

@@ -0,0 +1,52 @@
+#ifndef INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_
+#define INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_
+
+#include <cstdint>
+
+#include "native/cuda/caster.cuh"
+
+namespace infini::ops {
+
+template <Device::Type kDevice, typename Data, typename Index>
+__global__ void MoeSumKernel(
+    const Data* __restrict__ input, const Index* __restrict__ topk_ids,
+    const int32_t* __restrict__ expert_map, Data* __restrict__ output,
+    int64_t topk, int64_t hidden_size, int64_t input_token_stride,
+    int64_t input_slot_stride, int64_t input_hidden_stride,
+    int64_t topk_ids_token_stride, int64_t topk_ids_slot_stride,
+    int64_t expert_map_size, int64_t expert_map_stride) {
+  const int64_t token = static_cast<int64_t>(blockIdx.x);
+
+  for (int64_t hidden = static_cast<int64_t>(threadIdx.x); hidden < hidden_size;
+       hidden += static_cast<int64_t>(blockDim.x)) {
+    float sum = 0.0f;
+
+    for (int64_t slot = 0; slot < topk; ++slot) {
+      if (topk_ids != nullptr) {
+        const auto expert_id =
+            static_cast<int64_t>(topk_ids[token * topk_ids_token_stride +
+                                          slot * topk_ids_slot_stride]);
+        if (expert_id < 0) {
+          continue;
+        }
+        if (expert_map != nullptr &&
+            (expert_id >= expert_map_size ||
+             expert_map[expert_id * expert_map_stride] < 0)) {
+          continue;
+        }
+      }
+
+      const auto offset = token * input_token_stride +
+                          slot * input_slot_stride +
+                          hidden * input_hidden_stride;
+      sum += Caster<kDevice>::template Cast<float>(input[offset]);
+    }
+
+    output[token * hidden_size + hidden] =
+        Caster<kDevice>::template Cast<Data>(sum);
+  }
+}
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_