diff --git a/src/base/moe_align.h b/src/base/moe_align.h
new file mode 100644
index 000000000..f2857ea74
--- /dev/null
+++ b/src/base/moe_align.h
@@ -0,0 +1,71 @@
+#ifndef INFINI_OPS_BASE_MOE_ALIGN_H_
+#define INFINI_OPS_BASE_MOE_ALIGN_H_
+
+#include <cstdint>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class MoeAlign : public Operator<MoeAlign> {
+ public:
+  MoeAlign(Tensor sorted_token_ids, Tensor expert_ids,
+           Tensor num_tokens_post_padded, Tensor topk_ids, Tensor expert_map,
+           int64_t num_experts, int64_t block_size, bool pad_sorted_token_ids)
+      : sorted_token_ids_shape_{sorted_token_ids.shape()},
+        expert_ids_shape_{expert_ids.shape()},
+        num_tokens_post_padded_shape_{num_tokens_post_padded.shape()},
+        topk_ids_shape_{topk_ids.shape()},
+        expert_map_shape_{expert_map.shape()},
+        num_experts_{num_experts},
+        block_size_{block_size},
+        pad_sorted_token_ids_{pad_sorted_token_ids},
+        numel_{topk_ids.numel()},
+        max_num_tokens_padded_{sorted_token_ids.numel()} {
+    assert(topk_ids.ndim() == 2 && "`MoeAlign` topk_ids must be a 2D tensor");
+    assert(sorted_token_ids.ndim() == 1 &&
+           "`MoeAlign` sorted_token_ids must be a 1D tensor");
+    assert(expert_ids.ndim() == 1 &&
+           "`MoeAlign` expert_ids must be a 1D tensor");
+    assert(num_tokens_post_padded.ndim() == 1 &&
+           num_tokens_post_padded.numel() == 1 &&
+           "`MoeAlign` num_tokens_post_padded must be a scalar tensor");
+    assert(topk_ids.dtype() == DataType::kInt32 &&
+           "`MoeAlign` topk_ids must be int32");
+    assert(sorted_token_ids.dtype() == DataType::kInt32 &&
+           "`MoeAlign` sorted_token_ids must be int32");
+    assert(expert_ids.dtype() == DataType::kInt32 &&
+           "`MoeAlign` expert_ids must be int32");
+    assert(num_tokens_post_padded.dtype() == DataType::kInt32 &&
+           "`MoeAlign` num_tokens_post_padded must be int32");
+    assert(num_experts_ > 0 && "`MoeAlign` num_experts must be positive");
+    assert(block_size_ > 0 && "`MoeAlign` block_size must be positive");
+    assert(expert_map.numel() == 0 ||
+           (expert_map.ndim() == 1 &&
+            static_cast<int64_t>(expert_map.numel()) == num_experts_ &&
+            expert_map.dtype() == DataType::kInt32) &&
+               "`MoeAlign` expert_map must be empty or (num_experts,) int32");
+  }
+
+  virtual void operator()(Tensor sorted_token_ids, Tensor expert_ids,
+                          Tensor num_tokens_post_padded, Tensor topk_ids,
+                          Tensor expert_map, int64_t num_experts,
+                          int64_t block_size,
+                          bool pad_sorted_token_ids) const = 0;
+
+ protected:
+  Tensor::Shape sorted_token_ids_shape_;
+  Tensor::Shape expert_ids_shape_;
+  Tensor::Shape num_tokens_post_padded_shape_;
+  Tensor::Shape topk_ids_shape_;
+  Tensor::Shape expert_map_shape_;
+  int64_t num_experts_{0};
+  int64_t block_size_{0};
+  bool pad_sorted_token_ids_{false};
+  Tensor::Size numel_{0};
+  Tensor::Size max_num_tokens_padded_{0};
+};
+
+}  // namespace infini::ops
+
+#endif
diff --git a/src/native/cuda/metax/ops/moe_align/kernel.h b/src/native/cuda/metax/ops/moe_align/kernel.h
new file mode 100644
index 000000000..f40f6d40e
--- /dev/null
+++ b/src/native/cuda/metax/ops/moe_align/kernel.h
@@ -0,0 +1,20 @@
+#ifndef INFINI_OPS_METAX_MOE_ALIGN_KERNEL_H_
+#define INFINI_OPS_METAX_MOE_ALIGN_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/metax/runtime_.h"
+#include "native/cuda/ops/moe_align/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<MoeAlign, Device::Type::kMetax>
+    : public CudaMoeAlign<Runtime<Device::Type::kMetax>> {
+ public:
+  using CudaMoeAlign<Runtime<Device::Type::kMetax>>::CudaMoeAlign;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_METAX_MOE_ALIGN_KERNEL_H_
\ No newline at end of file
diff --git a/src/native/cuda/ops/moe_align/kernel.cuh b/src/native/cuda/ops/moe_align/kernel.cuh
new file mode 100644
index 000000000..574110314
--- /dev/null
+++ b/src/native/cuda/ops/moe_align/kernel.cuh
@@ -0,0 +1,280 @@
+/*
+ * Portions of the CUDA kernels in this file are adapted from SGLang:
+ * /sgl-kernel/csrc/moe/moe_align_kernel.cu
+ *
+ * Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#ifndef INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_CUH_
+#define INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_CUH_
+
+#include <cstddef>
+#include <cstdint>
+
+#include "native/cuda/kernel_commons.cuh"
+
+namespace infini::ops {
+
+namespace detail {
+
+constexpr int kMoeAlignVecSize = 4;
+using MoeAlignVec = int4;
+
+constexpr std::size_t MoeAlignNextPow2(std::size_t value) {
+  std::size_t result = 1;
+  while (result < value) {
+    result <<= 1;
+  }
+  return result;
+}
+
+template <typename T>
+constexpr T MoeAlignCeilDiv(T a, T b) {
+  return (a + b - 1) / b;
+}
+
+}  // namespace detail
+
+template <typename scalar_t>
+__global__ void MoeAlignCountAndSortExpertTokensKernel(
+    const int32_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ expert_map,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
+    size_t numel) {
+  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_map != nullptr) {
+      expert_id = expert_id >= 0 ? expert_map[expert_id] : -1;
+      if (expert_id < 0) {
+        continue;
+      }
+    }
+    expert_id += 1;
+    int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = static_cast<int32_t>(i);
+  }
+}
+
+template <typename scalar_t>
+__global__ void MoeAlignBlockSizeKernel(
+    const int32_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ expert_map,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, int32_t* __restrict__ cumsum,
+    bool pad_sorted_token_ids, const int32_t scan_size,
+    int32_t max_num_tokens_padded) {
+  if (blockIdx.x == 1) {
+    if (pad_sorted_token_ids) {
+      using detail::kMoeAlignVecSize;
+      using detail::MoeAlignVec;
+      MoeAlignVec fill_vec;
+      fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w =
+          static_cast<int32_t>(numel);
+      int32_t total_vecs =
+          (max_num_tokens_padded + kMoeAlignVecSize - 1) / kMoeAlignVecSize;
+      MoeAlignVec* out_ptr = reinterpret_cast<MoeAlignVec*>(sorted_token_ids);
+      for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
+        out_ptr[i] = fill_vec;
+      }
+    }
+    return;
+  }
+
+  extern __shared__ int32_t smem[];
+  int32_t* shared_counts = smem;
+  int32_t* prefix = shared_counts + num_experts;
+  int32_t* scan_buf = prefix + num_experts + 1;
+  __shared__ int32_t s_total_tokens_post_pad;
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+
+  if (tid < static_cast<size_t>(num_experts)) {
+    shared_counts[tid] = 0;
+  }
+
+  __syncthreads();
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_map != nullptr) {
+      expert_id = expert_id >= 0 ? expert_map[expert_id] : -1;
+      if (expert_id < 0) {
+        continue;
+      }
+    }
+    expert_id += 1;
+    atomicAdd(&shared_counts[expert_id], 1);
+  }
+
+  __syncthreads();
+
+  // scan_buf[i] = padded token count for expert slot i, zero elsewhere.
+  // NOTE: use a shared-memory (Hillis-Steele) scan instead of warp shuffles:
+  // on Metax hardware warpSize is 64 and __shfl_up_sync with a 32-bit mask
+  // does not propagate values across lane 32, silently corrupting the scan.
+  int32_t padded_count = 0;
+  if (tid < static_cast<size_t>(num_experts)) {
+    int32_t count = shared_counts[tid];
+    padded_count = (count + block_size - 1) / block_size * block_size;
+    scan_buf[tid] = padded_count;
+  }
+
+  // zero the tail [num_experts, scan_size) so the scan covers slots only.
+  if (tid >= static_cast<size_t>(num_experts) &&
+      tid < static_cast<size_t>(scan_size)) {
+    scan_buf[tid] = 0;
+  }
+  __syncthreads();
+
+  // block-wide inclusive scan (Hillis-Steele over shared memory;
+  // correct for any warp size, unlike a shuffle-based scan).
+  for (int32_t off = 1; off < scan_size; off <<= 1) {
+    int32_t self = 0, add = 0;
+    if (tid < static_cast<size_t>(scan_size)) {
+      self = scan_buf[tid];
+      if (tid >= static_cast<size_t>(off)) {
+        add = scan_buf[tid - off];
+      }
+    }
+    __syncthreads();
+    if (tid < static_cast<size_t>(scan_size)) {
+      scan_buf[tid] = self + add;
+    }
+    __syncthreads();
+  }
+
+  // exclusive prefix for slot i = inclusive_prefix[i] - padded_count[i]
+  if (tid < static_cast<size_t>(num_experts)) {
+    prefix[tid] = scan_buf[tid] - padded_count;
+  }
+  if (tid == 0) {
+    prefix[num_experts] = scan_buf[num_experts - 1];
+    s_total_tokens_post_pad = prefix[num_experts];
+    *total_tokens_post_pad = s_total_tokens_post_pad;
+  }
+  __syncthreads();
+
+  if (tid <= static_cast<size_t>(num_experts)) {
+    cumsum[tid] = prefix[tid];
+  }
+  __syncthreads();
+
+  const int32_t num_blocks = s_total_tokens_post_pad / block_size;
+  for (int32_t i = tid; i < num_blocks; i += stride) {
+    int32_t block_start = i * block_size;
+    int left = 0;
+    int right = num_experts;
+    while (left < right) {
+      int mid = (left + right) >> 1;
+      if (prefix[mid] <= block_start) {
+        left = mid + 1;
+      } else {
+        right = mid;
+      }
+    }
+    expert_ids[i] = left - 2;
+  }
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__global__ void MoeAlignBlockSizeSmallBatchExpertKernel(
+    const int32_t* __restrict__ topk_ids,
+    const int32_t* __restrict__ expert_map,
+    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, bool pad_sorted_token_ids,
+    int32_t max_num_tokens_padded) {
+  if (threadIdx.x < fill_threads) {
+    if (pad_sorted_token_ids) {
+      for (int32_t it = threadIdx.x; it < max_num_tokens_padded;
+           it += fill_threads) {
+        sorted_token_ids[it] = static_cast<int32_t>(numel);
+      }
+    }
+    __syncthreads();
+    __syncthreads();
+    __syncthreads();
+    return;
+  }
+
+  const size_t tid = threadIdx.x - fill_threads;
+  const size_t stride = blockDim.x - fill_threads;
+
+  extern __shared__ int32_t shared_mem[];
+  int32_t* cumsum = shared_mem;
+  int32_t* tokens_cnts =
+      reinterpret_cast<int32_t*>(shared_mem + num_experts + 1);
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[(tid + 1) * num_experts + i] = 0;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_map != nullptr) {
+      expert_id = expert_id >= 0 ? expert_map[expert_id] : -1;
+      if (expert_id < 0) {
+        continue;
+      }
+    }
+    expert_id += 1;
+    ++tokens_cnts[(tid + 1) * num_experts + expert_id];
+  }
+
+  __syncthreads();
+
+  if (tid < static_cast<size_t>(num_experts)) {
+    tokens_cnts[tid] = 0;
+    for (size_t i = 1; i <= stride; ++i) {
+      tokens_cnts[i * num_experts + tid] +=
+          tokens_cnts[(i - 1) * num_experts + tid];
+    }
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] = cumsum[i - 1] +
+                  detail::MoeAlignCeilDiv(
+                      tokens_cnts[stride * num_experts + i - 1], block_size) *
+                      block_size;
+    }
+    *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
+  }
+
+  __syncthreads();
+
+  if (tid < static_cast<size_t>(num_experts)) {
+    for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
+      expert_ids[i / block_size] = tid - 1;
+    }
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_map != nullptr) {
+      expert_id = expert_id >= 0 ? expert_map[expert_id] : -1;
+      if (expert_id < 0) {
+        continue;
+      }
+    }
+    expert_id += 1;
+    int32_t rank_post_pad =
+        tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+    sorted_token_ids[rank_post_pad] = static_cast<int32_t>(i);
+    ++tokens_cnts[tid * num_experts + expert_id];
+  }
+}
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_CUH_
\ No newline at end of file
diff --git a/src/native/cuda/ops/moe_align/kernel.h b/src/native/cuda/ops/moe_align/kernel.h
new file mode 100644
index 000000000..20460fc45
--- /dev/null
+++ b/src/native/cuda/ops/moe_align/kernel.h
@@ -0,0 +1,128 @@
+#ifndef INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_H_
+#define INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+
+#include "base/moe_align.h"
+#include "native/cuda/ops/moe_align/kernel.cuh"
+#include "native/cuda/runtime_utils.h"
+
+namespace infini::ops {
+
+namespace detail {
+
+constexpr int kMoeAlignWarpSize = 32;
+
+}  // namespace detail
+
+template <typename Backend>
+class CudaMoeAlign : public MoeAlign {
+ public:
+  CudaMoeAlign(Tensor sorted_token_ids, Tensor expert_ids,
+               Tensor num_tokens_post_padded, Tensor topk_ids,
+               Tensor expert_map, int64_t num_experts, int64_t block_size,
+               bool pad_sorted_token_ids)
+      : MoeAlign{sorted_token_ids, expert_ids,          num_tokens_post_padded,
+                 topk_ids,         expert_map,          num_experts,
+                 block_size,       pad_sorted_token_ids} {
+    // Scratch buffer for the per-expert cumulative (exclusive) prefix offsets.
+    const std::size_t cumsum_size =
+        (static_cast<std::size_t>(num_experts) + 1) * sizeof(int32_t);
+    Backend::Malloc(reinterpret_cast<void**>(&cumsum_buffer_), cumsum_size);
+  }
+
+  ~CudaMoeAlign() { Backend::Free(cumsum_buffer_); }
+
+  void operator()(Tensor sorted_token_ids, Tensor expert_ids,
+                  Tensor num_tokens_post_padded, Tensor topk_ids,
+                  Tensor expert_map, int64_t num_experts, int64_t block_size,
+                  bool pad_sorted_token_ids) const override {
+    int threads = RuntimeUtils<Backend::kDeviceType>::GetOptimalBlockSize();
+    threads = ((threads + detail::kMoeAlignWarpSize - 1) /
+               detail::kMoeAlignWarpSize) *
+              detail::kMoeAlignWarpSize;
+
+    const int32_t num_experts_shifted = static_cast<int32_t>(num_experts + 1);
+    const int32_t block_size_i32 = static_cast<int32_t>(block_size);
+    const int32_t max_num_tokens_padded =
+        static_cast<int32_t>(max_num_tokens_padded_);
+    const bool small_batch_expert_mode =
+        (numel_ < 1024) && (num_experts_shifted <= 64);
+
+    auto cuda_stream =
+        static_cast<typename Backend::Stream>(stream_ ? stream_ : 0);
+
+    const int32_t* d_topk_ids =
+        reinterpret_cast<const int32_t*>(topk_ids.data());
+    const int32_t* d_expert_map =
+        expert_map.numel() > 0
+            ? reinterpret_cast<const int32_t*>(expert_map.data())
+            : nullptr;
+    int32_t* d_sorted_token_ids =
+        reinterpret_cast<int32_t*>(sorted_token_ids.data());
+    int32_t* d_expert_ids = reinterpret_cast<int32_t*>(expert_ids.data());
+    int32_t* d_num_tokens_post_padded =
+        reinterpret_cast<int32_t*>(num_tokens_post_padded.data());
+
+    DispatchFunc<Backend::kDeviceType, DataType::kInt32>(
+        topk_ids.dtype(),
+        [&](auto type_tag) {
+          using T = typename decltype(type_tag)::type;
+
+          if (small_batch_expert_mode) {
+            constexpr int32_t fill_threads = 256;
+            const int32_t expert_threads =
+                std::max(num_experts_shifted, detail::kMoeAlignWarpSize);
+            const std::size_t shared_mem_size =
+                (static_cast<std::size_t>(expert_threads + 1) *
+                     static_cast<std::size_t>(num_experts_shifted) +
+                 static_cast<std::size_t>(num_experts_shifted + 1)) *
+                sizeof(int32_t);
+
+            MoeAlignBlockSizeSmallBatchExpertKernel<T, fill_threads>
+                <<<1, fill_threads + expert_threads, shared_mem_size,
+                   cuda_stream>>>(d_topk_ids, d_expert_map, d_sorted_token_ids,
+                                  d_expert_ids, d_num_tokens_post_padded,
+                                  num_experts_shifted, block_size_i32, numel_,
+                                  pad_sorted_token_ids, max_num_tokens_padded);
+          } else {
+            const int32_t scan_size = static_cast<int32_t>(
+                detail::MoeAlignNextPow2(num_experts_shifted));
+            const std::size_t shared_mem_size =
+                (static_cast<std::size_t>(num_experts_shifted) +
+                 static_cast<std::size_t>(num_experts_shifted + 1) +
+                 static_cast<std::size_t>(scan_size) +
+                 detail::kMoeAlignWarpSize) *
+                sizeof(int32_t);
+
+            MoeAlignBlockSizeKernel<T>
+                <<<2, threads, shared_mem_size, cuda_stream>>>(
+                    d_topk_ids, d_expert_map, d_sorted_token_ids, d_expert_ids,
+                    d_num_tokens_post_padded, num_experts_shifted,
+                    block_size_i32, numel_, cumsum_buffer_,
+                    pad_sorted_token_ids, scan_size, max_num_tokens_padded);
+
+            const int block_threads = std::min(256, threads);
+            const int num_blocks =
+                static_cast<int>((numel_ + block_threads - 1) / block_threads);
+            const int max_blocks = 65535;
+            const int actual_blocks = std::min(num_blocks, max_blocks);
+
+            MoeAlignCountAndSortExpertTokensKernel<T>
+                <<<actual_blocks, block_threads, 0, cuda_stream>>>(
+                    d_topk_ids, d_expert_map, d_sorted_token_ids,
+                    cumsum_buffer_, numel_);
+          }
+        },
+        "CudaMoeAlign::operator()");
+  }
+
+ private:
+  int32_t* cumsum_buffer_{nullptr};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CUDA_MOE_ALIGN_KERNEL_H_
\ No newline at end of file
diff --git a/tests/test_moe_align.py b/tests/test_moe_align.py
new file mode 100644
index 000000000..aeeb6e47c
--- /dev/null
+++ b/tests/test_moe_align.py
@@ -0,0 +1,122 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import infini.ops
+import pytest
+
+import torch
+from tests.utils import get_stream
+
+
+def _ref_moe_align(topk_ids, num_experts, block_size, pad_sorted_token_ids):
+    """Reference CPU implementation of moe_align."""
+    num_tokens, topk = topk_ids.shape
+    numel = num_tokens * topk
+
+    # Count tokens per expert
+    counts = [0] * (num_experts + 1)
+    for i in range(numel):
+        expert_id = topk_ids.view(-1)[i].item()
+        expert_id += 1  # Shift by 1
+        if 1 <= expert_id <= num_experts:
+            counts[expert_id] += 1
+
+    # Compute padded cumulative sums
+    prefix = [0] * (num_experts + 1)
+    total = 0
+    for i in range(1, num_experts + 1):
+        padded = (counts[i] + block_size - 1) // block_size * block_size
+        prefix[i] = total + padded
+        total = prefix[i]
+
+    num_tokens_post_padded = total
+    max_num_tokens_padded = prefix[num_experts]
+    max_num_blocks = (max_num_tokens_padded + block_size - 1) // block_size
+
+    # Fill sorted_token_ids
+    sorted_token_ids = torch.full(
+        (max_num_tokens_padded,), numel, dtype=torch.int32, device=topk_ids.device
+    )
+    positions = [0] * (num_experts + 1)  # Start from 0, not from counts[expert_id]
+    for i in range(numel):
+        expert_id = topk_ids.view(-1)[i].item()
+        expert_id += 1
+        if 1 <= expert_id <= num_experts:
+            rank = prefix[expert_id - 1] + positions[expert_id]
+            sorted_token_ids[rank] = i
+            positions[expert_id] += 1
+
+    # Fill expert_ids
+    expert_ids = torch.full(
+        (max_num_blocks,), -1, dtype=torch.int32, device=topk_ids.device
+    )
+    num_blocks = num_tokens_post_padded // block_size
+    for i in range(num_blocks):
+        block_start = i * block_size
+        expert = 0
+        for e in range(1, num_experts + 1):
+            if prefix[e] > block_start:
+                expert = e - 1
+                break
+        expert_ids[i] = expert
+
+    return sorted_token_ids, expert_ids, torch.tensor(
+        [num_tokens_post_padded], dtype=torch.int32, device=topk_ids.device
+    )
+
+
+@pytest.mark.parametrize(
+    "num_tokens, topk, num_experts, block_size, pad_sorted_token_ids",
+    (
+        (4, 2, 2, 4, True),
+        (8, 2, 4, 4, True),
+        (8, 2, 4, 4, False),
+        (16, 2, 8, 8, True),
+        (1, 1, 2, 4, True),
+    ),
+)
+def test_moe_align(
+    num_tokens,
+    topk,
+    num_experts,
+    block_size,
+    pad_sorted_token_ids,
+    device,
+):
+    topk_ids = torch.randint(
+        0, num_experts, (num_tokens, topk), dtype=torch.int32, device=device
+    )
+
+    ref_sorted, ref_expert_ids, ref_num_tokens = _ref_moe_align(
+        topk_ids, num_experts, block_size, pad_sorted_token_ids
+    )
+
+    # Kernel skips padding init when pad_sorted_token_ids=False; pre-fill
+    # with the sentinel so undefined positions match the reference.
+    sorted_token_ids = (
+        torch.full_like(ref_sorted, num_tokens * topk)
+        if not pad_sorted_token_ids
+        else torch.empty_like(ref_sorted)
+    )
+    expert_ids = torch.empty_like(ref_expert_ids)
+    num_tokens_post_padded = torch.empty_like(ref_num_tokens)
+
+    expert_map = torch.empty(0, dtype=torch.int32, device=device)
+
+    infini.ops.moe_align(
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        topk_ids,
+        expert_map,
+        num_experts,
+        block_size,
+        pad_sorted_token_ids,
+        stream=get_stream(device),
+    )
+
+    torch.testing.assert_close(sorted_token_ids, ref_sorted)
+    torch.testing.assert_close(expert_ids, ref_expert_ids)
+    torch.testing.assert_close(num_tokens_post_padded, ref_num_tokens)