feat(metax): add moe_fused_dense operator for fused MoE expert forward

src/base/moe_fused_dense.h

@@ -0,0 +1,99 @@
+#ifndef INFINI_OPS_BASE_MOE_FUSED_DENSE_H_
+#define INFINI_OPS_BASE_MOE_FUSED_DENSE_H_
+
+#include <cstdint>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class MoeFusedDense : public Operator<MoeFusedDense> {
+ public:
+  MoeFusedDense(Tensor output, Tensor hidden_states, Tensor w13, Tensor w2,
+                Tensor topk_weights, Tensor topk_ids, Tensor sorted_token_ids,
+                Tensor expert_ids, Tensor num_tokens_post_padded)
+      : output_shape_{output.shape()},
+        hidden_states_shape_{hidden_states.shape()},
+        w13_shape_{w13.shape()},
+        w2_shape_{w2.shape()},
+        topk_weights_shape_{topk_weights.shape()},
+        topk_ids_shape_{topk_ids.shape()},
+        sorted_token_ids_shape_{sorted_token_ids.shape()},
+        expert_ids_shape_{expert_ids.shape()},
+        num_tokens_post_padded_shape_{num_tokens_post_padded.shape()},
+        num_tokens_{hidden_states.size(0)},
+        hidden_size_{hidden_states.size(1)},
+        num_experts_{w13.size(0)},
+        intermediate_size_{w2.size(2)},
+        topk_{static_cast<int64_t>(topk_ids.size(1))},
+        max_num_tokens_padded_{sorted_token_ids.numel()},
+        max_num_blocks_{expert_ids.numel()},
+        dtype_{output.dtype()} {
+    assert(output.ndim() == 2 && hidden_states.ndim() == 2 &&
+           "`MoeFusedDense` output and hidden_states must be 2D tensors");
+    assert(w13.ndim() == 3 && w2.ndim() == 3 &&
+           "`MoeFusedDense` w13 and w2 must be 3D tensors");
+    assert(topk_weights.ndim() == 2 && topk_ids.ndim() == 2 &&
+           "`MoeFusedDense` topk_weights and topk_ids must be 2D tensors");
+    assert(
+        sorted_token_ids.ndim() == 1 && expert_ids.ndim() == 1 &&
+        "`MoeFusedDense` sorted_token_ids and expert_ids must be 1D tensors");
+    assert(num_tokens_post_padded.ndim() == 1 &&
+           num_tokens_post_padded.numel() == 1 &&
+           "`MoeFusedDense` num_tokens_post_padded must be a scalar tensor");
+    assert(
+        output.dtype() == hidden_states.dtype() &&
+        output.dtype() == w13.dtype() && output.dtype() == w2.dtype() &&
+        "`MoeFusedDense` all weight/output tensors must have the same dtype");
+    assert(topk_weights.dtype() == DataType::kFloat32 &&
+           "`MoeFusedDense` topk_weights must be float32");
+    assert(topk_ids.dtype() == DataType::kInt32 &&
+           sorted_token_ids.dtype() == DataType::kInt32 &&
+           expert_ids.dtype() == DataType::kInt32 &&
+           num_tokens_post_padded.dtype() == DataType::kInt32 &&
+           "`MoeFusedDense` index tensors must be int32");
+    assert(output.size(0) == num_tokens_ && output.size(1) == hidden_size_ &&
+           "`MoeFusedDense` output shape must be (num_tokens, hidden_size)");
+    assert(w13.size(2) == hidden_size_ &&
+           "`MoeFusedDense` w13 must have shape (num_experts, w13_rows, "
+           "hidden_size)");
+    assert(w2.size(0) == num_experts_ && w2.size(1) == hidden_size_ &&
+           "`MoeFusedDense` w2 must have shape (num_experts, hidden_size, "
+           "intermediate_size)");
+    assert(topk_weights.size(0) == num_tokens_ &&
+           topk_weights.size(1) == static_cast<Tensor::Size>(topk_) &&
+           topk_ids.size(0) == num_tokens_ &&
+           "`MoeFusedDense` topk shapes must match");
+    assert(max_num_tokens_padded_ >= num_tokens_ * topk_ &&
+           max_num_blocks_ > 0 &&
+           "`MoeFusedDense` invalid sorted_token_ids or expert_ids sizes");
+  }
+
+  virtual void operator()(Tensor output, Tensor hidden_states, Tensor w13,
+                          Tensor w2, Tensor topk_weights, Tensor topk_ids,
+                          Tensor sorted_token_ids, Tensor expert_ids,
+                          Tensor num_tokens_post_padded) const = 0;
+
+ protected:
+  Tensor::Shape output_shape_;
+  Tensor::Shape hidden_states_shape_;
+  Tensor::Shape w13_shape_;
+  Tensor::Shape w2_shape_;
+  Tensor::Shape topk_weights_shape_;
+  Tensor::Shape topk_ids_shape_;
+  Tensor::Shape sorted_token_ids_shape_;
+  Tensor::Shape expert_ids_shape_;
+  Tensor::Shape num_tokens_post_padded_shape_;
+  Tensor::Size num_tokens_{0};
+  Tensor::Size hidden_size_{0};
+  Tensor::Size num_experts_{0};
+  Tensor::Size intermediate_size_{0};
+  int64_t topk_{0};
+  Tensor::Size max_num_tokens_padded_{0};
+  Tensor::Size max_num_blocks_{0};
+  DataType dtype_;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/metax/ops/moe_fused_dense/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_
+#define INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/metax/blas.h"
+#include "native/cuda/metax/runtime_.h"
+#include "native/cuda/ops/moe_fused_dense/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<MoeFusedDense, Device::Type::kMetax>
+    : public CudaMoeFusedDense<Blas<Device::Type::kMetax>> {
+ public:
+  using CudaMoeFusedDense<Blas<Device::Type::kMetax>>::CudaMoeFusedDense;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_

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

@@ -0,0 +1,130 @@
+#ifndef INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_
+#define INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_
+
+#include <cmath>
+
+#include "native/cuda/kernel_commons.cuh"
+
+namespace infini::ops {
+
+// Computes the inclusive prefix sum of `counts[0..num_experts-1]`, storing the
+// exclusive prefix in `offsets[0..num_experts-1]` and the total in
+// `offsets[num_experts]`. Only one thread (thread 0) is used.
+__global__ void ExclusivePrefixCountsKernel(const int* counts, int* offsets,
+                                            int num_experts) {
+  if (threadIdx.x == 0) {
+    offsets[0] = 0;
+    int sum = 0;
+    for (int i = 0; i < num_experts; ++i) {
+      sum += counts[i];
+      offsets[i + 1] = sum;
+    }
+  }
+}
+
+// For each aligned block of `block_size` tokens, atomically adds `block_size`
+// to the count of the expert that owns the block. Rows whose expert id is out
+// of range (padding rows) are skipped.
+// NOTE: `num_tokens_post_padded` is a host-side scalar (not a device pointer),
+// because Metax does not support device-side pointer dereference inside
+// kernels.
+__global__ void CountAlignedExpertsKernel(const int* expert_ids,
+                                          int num_tokens_post_padded,
+                                          int* counts, int num_experts,
+                                          int block_size) {
+  int block = blockIdx.x * blockDim.x + threadIdx.x;
+  int num_blocks = (num_tokens_post_padded + block_size - 1) / block_size;
+  if (block >= num_blocks) {
+    return;
+  }
+  int expert = expert_ids[block];
+  if (expert >= 0 && expert < num_experts) {
+    atomicAdd(counts + expert, block_size);
+  }
+}
+
+// Gathers hidden states into a packed, expert-bucketed buffer. Each output row
+// `row` maps to `pair = sorted_token_ids[row]`, and its hidden state is copied
+// from the source token `pair / topk`. Rows pointing past the valid pair range
+// (padding) are zero-filled.
+template <Device::Type kDev, typename T>
+__global__ void PackHiddenAlignedKernel(const T* hidden,
+                                        const int* sorted_token_ids,
+                                        int* output_permutation,
+                                        T* packed_hidden, int pairs, int topk,
+                                        int hidden_size,
+                                        int max_num_tokens_padded) {
+  int row = blockIdx.x;
+  int tid = threadIdx.x;
+  if (row >= max_num_tokens_padded) {
+    return;
+  }
+  int pair = sorted_token_ids[row];
+  if (pair >= 0 && pair < pairs) {
+    if (tid == 0) {
+      output_permutation[pair] = row;
+    }
+    int token = pair / topk;
+    for (int h = tid; h < hidden_size; h += blockDim.x) {
+      packed_hidden[static_cast<size_t>(row) * hidden_size + h] =
+          hidden[static_cast<size_t>(token) * hidden_size + h];
+    }
+  } else {
+    for (int h = tid; h < hidden_size; h += blockDim.x) {
+      packed_hidden[static_cast<size_t>(row) * hidden_size + h] =
+          Caster<kDev>::template Cast<T>(0.0f);
+    }
+  }
+}
+
+// SwiGLU activation: out = up * silu(gate), where silu(x) = x / (1 + exp(-x)).
+template <Device::Type kDev, typename T>
+__global__ void SwigluKernel(const T* gate_up, T* activated, int rows,
+                             int intermediate_size) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int total = rows * intermediate_size;
+  if (idx >= total) {
+    return;
+  }
+  int row = idx / intermediate_size;
+  int col = idx - row * intermediate_size;
+  const T* base = gate_up + static_cast<size_t>(row) * intermediate_size * 2;
+  float gate = Caster<kDev>::template Cast<float>(base[col]);
+  float up = Caster<kDev>::template Cast<float>(base[intermediate_size + col]);
+  float silu = gate / (1.0f + expf(-gate));
+  activated[idx] = Caster<kDev>::template Cast<T>(up * silu);
+}
+
+// Scatters the weighted expert outputs back to the original token rows. For
+// each token `token`, sums over its `topk` pairs, gathering rows through
+// `output_permutation` and weighting by `topk_weights`.
+template <Device::Type kDev, typename T>
+__global__ void ApplyShuffleMulSumKernel(
+    const T* __restrict__ expert_out, T* __restrict__ out,
+    const int* __restrict__ output_permutation,
+    const float* __restrict__ topk_weights, 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 += Caster<kDev>::template Cast<float>(
+                   expert_out[static_cast<size_t>(src_row) * hidden_size + h]) *
+               topk_weights[pair];
+      }
+    }
+    out[static_cast<size_t>(token) * hidden_size + h] =
+        Caster<kDev>::template Cast<T>(sum);
+  }
+}
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_