diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
index 1fa2c0d18e7..152bc116edb 100644
--- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
@@ -82,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel(
 }
 }  // namespace batched_moe_align_block_size
 
+template <typename scalar_t>
+__device__ __forceinline__ int get_local_expert_id(
+    size_t idx, const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ expert_map, int32_t num_experts,
+    bool has_expert_map) {
+  int expert_id = topk_ids[idx];
+  if (expert_id >= num_experts || expert_id < 0) {
+    return -1;
+  }
+  if (has_expert_map) {
+    expert_id = expert_map[expert_id];
+  }
+  return expert_id;
+}
+
 template <typename scalar_t>
 __device__ void _moe_align_block_size(
     const scalar_t* __restrict__ topk_ids,
@@ -126,20 +141,15 @@ __device__ void _moe_align_block_size(
   const size_t stride = blockDim.x;
 
   for (size_t i = tid; i < numel; i += stride) {
-    int expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int warp_idx = expert_id / experts_per_warp;
+      int expert_offset = expert_id % experts_per_warp;
+      int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+      atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
+                mask);
     }
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid experts
-      if (expert_id == -1) continue;
-    }
-    int warp_idx = expert_id / experts_per_warp;
-    int expert_offset = expert_id % experts_per_warp;
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
-              mask);
   }
 
   __syncthreads();
@@ -227,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert(
   }
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid expert
-      if (expert_id == -1) continue;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+      tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
     }
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
   }
 
   __syncthreads();
@@ -276,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert(
   }
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid expert
-      if (expert_id == -1) continue;
-    }
-    int32_t rank_post_pad =
-        tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int32_t rank_post_pad =
+          tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
 
-    if (token_mask == nullptr || token_mask[i / topk_num]) {
-      sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
-      ++tokens_cnts[tid * num_experts + expert_id];
+      if (token_mask == nullptr || token_mask[i / topk_num]) {
+        sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
+        ++tokens_cnts[tid * num_experts + expert_id];
+      }
     }
   }
 }
@@ -303,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens(
   const size_t stride = blockDim.x * gridDim.y;
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
-    }
-
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid experts
-      if (expert_id == -1) continue;
-    }
-
-    if (token_mask == nullptr || token_mask[i / topk_num]) {
-      int32_t rank_post_pad = atomicAdd(
-          &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
-      sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
-          i;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      if (token_mask == nullptr || token_mask[i / topk_num]) {
+        int32_t rank_post_pad = atomicAdd(
+            &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
+        sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
+            i;
+      }
     }
   }
 }
diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py
index 9096d0ab856..a017fa07bf9 100644
--- a/tests/kernels/moe/test_moe_align_block_size.py
+++ b/tests/kernels/moe/test_moe_align_block_size.py
@@ -246,20 +246,30 @@ def test_moe_align_block_size(
 @pytest.mark.parametrize("topk", [2, 4])
 @pytest.mark.parametrize("num_experts", [8, 64])
 @pytest.mark.parametrize("block_size", [64])
+@pytest.mark.parametrize("mask_inactive_experts", [False, True])
 def test_moe_align_block_size_with_expert_map(
-    m: int, topk: int, num_experts: int, block_size: int
+    m: int,
+    topk: int,
+    num_experts: int,
+    block_size: int,
+    mask_inactive_experts: bool,
 ):
     """Test moe_align_block_size with expert mapping (EP scenario)"""
-    topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32)
-    for i in range(m):
-        experts = torch.randperm(num_experts, device="cuda")[:topk]
-        topk_ids[i] = experts
-
     expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
     local_experts = list(range(0, num_experts, 2))
     for i, expert_id in enumerate(local_experts):
         expert_map[expert_id] = i
 
+    topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32)
+    for i in range(m):
+        experts = torch.randperm(num_experts, device="cuda")[:topk]
+        for k in range(topk):
+            topk_ids[i, k] = (
+                experts[k]
+                if (experts[k] in local_experts) or not mask_inactive_experts
+                else -1
+            )
+
     actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size(
         topk_ids=topk_ids,
         block_size=block_size,