feat(nvidia): add moe_align_block_size operator

src/base/moe_align_block_size.h

@@ -0,0 +1,207 @@
+#ifndef INFINI_OPS_BASE_MOE_ALIGN_BLOCK_SIZE_H_
+#define INFINI_OPS_BASE_MOE_ALIGN_BLOCK_SIZE_H_
+
+#include <cassert>
+#include <cstdint>
+#include <functional>
+#include <limits>
+#include <optional>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Align routed token indices into expert-specific blocks following vLLM's
+// low-level `moe_align_block_size` operator.
+class MoeAlignBlockSize : public Operator<MoeAlignBlockSize> {
+ public:
+  MoeAlignBlockSize(const Tensor topk_ids, const int64_t num_experts,
+                    const int64_t block_size, Tensor sorted_token_ids,
+                    Tensor experts_ids, Tensor num_tokens_post_pad)
+      : topk_ids_metadata_{topk_ids},
+        expert_map_metadata_{std::nullopt},
+        sorted_token_ids_metadata_{sorted_token_ids},
+        experts_ids_metadata_{experts_ids},
+        num_tokens_post_pad_metadata_{num_tokens_post_pad},
+        numel_{topk_ids.numel()},
+        num_experts_{num_experts},
+        block_size_{block_size},
+        sorted_token_ids_size_{sorted_token_ids.numel()},
+        experts_ids_size_{experts_ids.numel()} {
+    Validate(topk_ids, std::nullopt, sorted_token_ids, experts_ids,
+             num_tokens_post_pad);
+  }
+
+  MoeAlignBlockSize(const Tensor topk_ids, const Tensor expert_map,
+                    const int64_t num_experts, const int64_t block_size,
+                    Tensor sorted_token_ids, Tensor experts_ids,
+                    Tensor num_tokens_post_pad)
+      : topk_ids_metadata_{topk_ids},
+        expert_map_metadata_{expert_map},
+        sorted_token_ids_metadata_{sorted_token_ids},
+        experts_ids_metadata_{experts_ids},
+        num_tokens_post_pad_metadata_{num_tokens_post_pad},
+        numel_{topk_ids.numel()},
+        num_experts_{num_experts},
+        block_size_{block_size},
+        sorted_token_ids_size_{sorted_token_ids.numel()},
+        experts_ids_size_{experts_ids.numel()} {
+    Validate(topk_ids, std::optional<Tensor>{expert_map}, sorted_token_ids,
+             experts_ids, num_tokens_post_pad);
+  }
+
+  void operator()(const Tensor topk_ids, const int64_t num_experts,
+                  const int64_t block_size, Tensor sorted_token_ids,
+                  Tensor experts_ids, Tensor num_tokens_post_pad) const {
+    ValidateInvocation(topk_ids, std::nullopt, num_experts, block_size,
+                       sorted_token_ids, experts_ids, num_tokens_post_pad);
+    Run(topk_ids, std::nullopt, num_experts, block_size, sorted_token_ids,
+        experts_ids, num_tokens_post_pad);
+  }
+
+  void operator()(const Tensor topk_ids, const Tensor expert_map,
+                  const int64_t num_experts, const int64_t block_size,
+                  Tensor sorted_token_ids, Tensor experts_ids,
+                  Tensor num_tokens_post_pad) const {
+    ValidateInvocation(topk_ids, std::optional<Tensor>{expert_map}, num_experts,
+                       block_size, sorted_token_ids, experts_ids,
+                       num_tokens_post_pad);
+    Run(topk_ids, std::optional<Tensor>{expert_map}, num_experts, block_size,
+        sorted_token_ids, experts_ids, num_tokens_post_pad);
+  }
+
+ protected:
+  void ValidateInvocation(const Tensor topk_ids,
+                          const std::optional<Tensor> maybe_expert_map,
+                          const int64_t num_experts, const int64_t block_size,
+                          Tensor sorted_token_ids, Tensor experts_ids,
+                          Tensor num_tokens_post_pad) const {
+    assert(num_experts == num_experts_ && block_size == block_size_ &&
+           "`MoeAlignBlockSize` attributes changed after descriptor creation");
+
+    assert(CallMetadataMatches(topk_ids, maybe_expert_map, sorted_token_ids,
+                               experts_ids, num_tokens_post_pad) &&
+           "`MoeAlignBlockSize` tensor metadata differs from its descriptor");
+  }
+
+  bool CallMetadataMatches(const Tensor topk_ids,
+                           const std::optional<Tensor> maybe_expert_map,
+                           const Tensor sorted_token_ids,
+                           const Tensor experts_ids,
+                           const Tensor num_tokens_post_pad) const {
+    const std::equal_to<Tensor> same_metadata;
+    const auto same_expert_map_metadata =
+        expert_map_metadata_.has_value() == maybe_expert_map.has_value() &&
+        (!expert_map_metadata_ ||
+         same_metadata(*expert_map_metadata_, *maybe_expert_map));
+
+    return same_metadata(topk_ids_metadata_, topk_ids) &&
+           same_expert_map_metadata &&
+           same_metadata(sorted_token_ids_metadata_, sorted_token_ids) &&
+           same_metadata(experts_ids_metadata_, experts_ids) &&
+           same_metadata(num_tokens_post_pad_metadata_, num_tokens_post_pad);
+  }
+
+  void Validate(const Tensor topk_ids,
+                const std::optional<Tensor> maybe_expert_map,
+                Tensor sorted_token_ids, Tensor experts_ids,
+                Tensor num_tokens_post_pad) const {
+    assert(topk_ids.ndim() == 2 &&
+           "`MoeAlignBlockSize` requires 2D `topk_ids`");
+    assert(topk_ids.dtype() == DataType::kInt32 &&
+           "`MoeAlignBlockSize` currently requires int32 `topk_ids`");
+    assert(topk_ids.IsContiguous() &&
+           "`MoeAlignBlockSize` requires contiguous `topk_ids`");
+    assert(num_experts_ > 0 && num_experts_ < 1024 &&
+           "`MoeAlignBlockSize` requires `num_experts` in [1, 1023]");
+    assert(block_size_ > 0 &&
+           "`MoeAlignBlockSize` requires a positive `block_size`");
+    assert(block_size_ <= std::numeric_limits<int32_t>::max() &&
+           numel_ <=
+               static_cast<Tensor::Size>(std::numeric_limits<int32_t>::max()) &&
+           "`MoeAlignBlockSize` requires int32-addressable token indices");
+
+    const auto same_device_as_topk_ids = [&](const Tensor tensor) {
+      return tensor.device().type() == topk_ids.device().type() &&
+             tensor.device().index() == topk_ids.device().index();
+    };
+    assert(same_device_as_topk_ids(sorted_token_ids) &&
+           same_device_as_topk_ids(experts_ids) &&
+           same_device_as_topk_ids(num_tokens_post_pad) &&
+           "`MoeAlignBlockSize` requires all tensors on the same device");
+
+    assert(sorted_token_ids.ndim() == 1 && experts_ids.ndim() == 1 &&
+           num_tokens_post_pad.ndim() == 1 &&
+           num_tokens_post_pad.numel() == 1 &&
+           "`MoeAlignBlockSize` requires 1D output tensors");
+    assert(sorted_token_ids.dtype() == DataType::kInt32 &&
+           experts_ids.dtype() == DataType::kInt32 &&
+           num_tokens_post_pad.dtype() == DataType::kInt32 &&
+           "`MoeAlignBlockSize` requires int32 output tensors");
+    assert(sorted_token_ids.IsContiguous() && experts_ids.IsContiguous() &&
+           num_tokens_post_pad.IsContiguous() &&
+           "`MoeAlignBlockSize` requires contiguous output tensors");
+
+    const auto num_experts = static_cast<Tensor::Size>(num_experts_);
+    const auto block_size = static_cast<Tensor::Size>(block_size_);
+    auto required_sorted_size = numel_ + num_experts * (block_size - 1);
+    if (numel_ < num_experts) {
+      const auto small_input_size = numel_ * block_size;
+      required_sorted_size = small_input_size < required_sorted_size
+                                 ? small_input_size
+                                 : required_sorted_size;
+    }
+    const auto required_experts_size =
+        (required_sorted_size + block_size - 1) / block_size;
+    assert(required_sorted_size <=
+               static_cast<Tensor::Size>(std::numeric_limits<int32_t>::max()) &&
+           "`MoeAlignBlockSize` requires int32-addressable padded indices");
+    assert(sorted_token_ids_size_ >= required_sorted_size &&
+           experts_ids_size_ >= required_experts_size &&
+           "`MoeAlignBlockSize` output tensors are too small");
+
+    if (maybe_expert_map) {
+      assert(maybe_expert_map->ndim() == 1 &&
+             maybe_expert_map->numel() ==
+                 static_cast<Tensor::Size>(num_experts_) &&
+             "`MoeAlignBlockSize` requires `expert_map` shape "
+             "[`num_experts`]");
+      assert(maybe_expert_map->dtype() == DataType::kInt32 &&
+             "`MoeAlignBlockSize` currently requires int32 `expert_map`");
+      assert(maybe_expert_map->IsContiguous() &&
+             "`MoeAlignBlockSize` requires contiguous `expert_map`");
+      assert(same_device_as_topk_ids(*maybe_expert_map) &&
+             "`MoeAlignBlockSize` requires `expert_map` on the input device");
+    }
+  }
+
+  virtual void Run(const Tensor topk_ids,
+                   const std::optional<Tensor> maybe_expert_map,
+                   const int64_t num_experts, const int64_t block_size,
+                   Tensor sorted_token_ids, Tensor experts_ids,
+                   Tensor num_tokens_post_pad) const = 0;
+
+  Tensor topk_ids_metadata_;
+
+  std::optional<Tensor> expert_map_metadata_;
+
+  Tensor sorted_token_ids_metadata_;
+
+  Tensor experts_ids_metadata_;
+
+  Tensor num_tokens_post_pad_metadata_;
+
+  Tensor::Size numel_{0};
+
+  int64_t num_experts_{0};
+
+  int64_t block_size_{0};
+
+  Tensor::Size sorted_token_ids_size_{0};
+
+  Tensor::Size experts_ids_size_{0};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_BASE_MOE_ALIGN_BLOCK_SIZE_H_

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

@@ -0,0 +1,87 @@
+#include <cuda_runtime_api.h>
+
+#include <cassert>
+#include <cstdint>
+#include <optional>
+
+#include "native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh"
+#include "native/cuda/nvidia/ops/moe_align_block_size/kernel.h"
+
+namespace infini::ops {
+namespace {
+
+class DeviceGuard {
+ public:
+  explicit DeviceGuard(int device_index) {
+    auto status = cudaGetDevice(&previous_device_);
+    assert(status == cudaSuccess &&
+           "`MoeAlignBlockSize` failed to query the current CUDA device");
+
+    if (previous_device_ != device_index) {
+      status = cudaSetDevice(device_index);
+      assert(status == cudaSuccess &&
+             "`MoeAlignBlockSize` failed to select the input CUDA device");
+      restore_ = true;
+    }
+  }
+
+  ~DeviceGuard() {
+    if (restore_) {
+      const auto status = cudaSetDevice(previous_device_);
+      assert(status == cudaSuccess &&
+             "`MoeAlignBlockSize` failed to restore the CUDA device");
+    }
+  }
+
+ private:
+  int previous_device_{0};
+
+  bool restore_{false};
+};
+
+}  // namespace
+
+Operator<MoeAlignBlockSize, Device::Type::kNvidia, 0>::Operator(
+    const Tensor topk_ids, const int64_t num_experts, const int64_t block_size,
+    Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad)
+    : MoeAlignBlockSize{topk_ids,         num_experts, block_size,
+                        sorted_token_ids, experts_ids, num_tokens_post_pad},
+      device_index_{topk_ids.device().index()} {}
+
+Operator<MoeAlignBlockSize, Device::Type::kNvidia, 0>::Operator(
+    const Tensor topk_ids, const Tensor expert_map, const int64_t num_experts,
+    const int64_t block_size, Tensor sorted_token_ids, Tensor experts_ids,
+    Tensor num_tokens_post_pad)
+    : MoeAlignBlockSize{topk_ids,           expert_map,       num_experts,
+                        block_size,         sorted_token_ids, experts_ids,
+                        num_tokens_post_pad},
+      device_index_{topk_ids.device().index()} {}
+
+void Operator<MoeAlignBlockSize, Device::Type::kNvidia, 0>::Run(
+    const Tensor topk_ids, const std::optional<Tensor> maybe_expert_map,
+    const int64_t num_experts, const int64_t block_size,
+    Tensor sorted_token_ids, Tensor experts_ids,
+    Tensor num_tokens_post_pad) const {
+  DeviceGuard device_guard{device_index_};
+
+  constexpr int32_t kThreads = 256;
+  const auto shared_memory_size =
+      static_cast<size_t>(num_experts_) * sizeof(int32_t);
+  moe_align_block_size_detail::MoeAlignBlockSizeKernel<<<
+      1, kThreads, shared_memory_size, static_cast<cudaStream_t>(stream_)>>>(
+      reinterpret_cast<const int32_t*>(topk_ids.data()),
+      maybe_expert_map
+          ? reinterpret_cast<const int32_t*>(maybe_expert_map->data())
+          : nullptr,
+      reinterpret_cast<int32_t*>(sorted_token_ids.data()),
+      reinterpret_cast<int32_t*>(experts_ids.data()),
+      reinterpret_cast<int32_t*>(num_tokens_post_pad.data()), numel_,
+      static_cast<int32_t>(num_experts_), static_cast<int32_t>(block_size_),
+      sorted_token_ids_size_, experts_ids_size_);
+
+  const auto status = cudaGetLastError();
+  assert(status == cudaSuccess &&
+         "`MoeAlignBlockSize` CUDA kernel launch failed");
+}
+
+}  // namespace infini::ops

src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh

@@ -0,0 +1,83 @@
+#ifndef INFINI_OPS_NVIDIA_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_
+#define INFINI_OPS_NVIDIA_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_
+
+#include <cstddef>
+#include <cstdint>
+
+namespace infini::ops {
+namespace moe_align_block_size_detail {
+
+__device__ __forceinline__ int32_t MapExpertId(int32_t expert_id,
+                                               const int32_t* expert_map,
+                                               int32_t num_experts) {
+  if (expert_id < 0 || expert_id >= num_experts) {
+    return -1;
+  }
+  if (expert_map != nullptr) {
+    expert_id = expert_map[expert_id];
+  }
+
+  return expert_id >= 0 && expert_id < num_experts ? expert_id : -1;
+}
+
+__global__ void MoeAlignBlockSizeKernel(
+    const int32_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ expert_map,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ experts_ids,
+    int32_t* __restrict__ num_tokens_post_pad, size_t numel,
+    int32_t num_experts, int32_t block_size, size_t sorted_token_ids_size,
+    size_t experts_ids_size) {
+  extern __shared__ int32_t expert_offsets[];
+
+  for (size_t i = threadIdx.x; i < sorted_token_ids_size; i += blockDim.x) {
+    sorted_token_ids[i] = static_cast<int32_t>(numel);
+  }
+  for (size_t i = threadIdx.x; i < experts_ids_size; i += blockDim.x) {
+    experts_ids[i] = -1;
+  }
+  for (int32_t i = threadIdx.x; i < num_experts; i += blockDim.x) {
+    expert_offsets[i] = 0;
+  }
+  if (threadIdx.x == 0) {
+    *num_tokens_post_pad = 0;
+  }
+  __syncthreads();
+
+  for (size_t i = threadIdx.x; i < numel; i += blockDim.x) {
+    const auto expert_id = MapExpertId(topk_ids[i], expert_map, num_experts);
+    if (expert_id >= 0) {
+      atomicAdd(&expert_offsets[expert_id], 1);
+    }
+  }
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    int32_t token_offset = 0;
+    for (int32_t expert_id = 0; expert_id < num_experts; ++expert_id) {
+      const int32_t count = expert_offsets[expert_id];
+      expert_offsets[expert_id] = token_offset;
+      const int32_t padded_count =
+          (count + block_size - 1) / block_size * block_size;
+      for (int32_t i = 0; i < padded_count; i += block_size) {
+        experts_ids[(token_offset + i) / block_size] = expert_id;
+      }
+      token_offset += padded_count;
+    }
+    *num_tokens_post_pad = token_offset;
+  }
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    for (size_t i = 0; i < numel; ++i) {
+      const auto expert_id = MapExpertId(topk_ids[i], expert_map, num_experts);
+      if (expert_id >= 0) {
+        sorted_token_ids[expert_offsets[expert_id]++] = static_cast<int32_t>(i);
+      }
+    }
+  }
+}
+
+}  // namespace moe_align_block_size_detail
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_NVIDIA_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_