diff --git a/tile_kernels/moe/expand_to_fused_kernel.py b/tile_kernels/moe/expand_to_fused_kernel.py
index 92a18c4..87bee96 100644
--- a/tile_kernels/moe/expand_to_fused_kernel.py
+++ b/tile_kernels/moe/expand_to_fused_kernel.py
@@ -56,13 +56,15 @@ def expand_to_fused_kernel(
             if pid_token < num_expanded_tokens:
                 if pos_to_expert[pid_token] < 0:
                     for i in T.Parallel(hidden_aligned):
-                        expanded_x[pid_token, i] = 0
+                        if i < hidden:
+                            expanded_x[pid_token, i] = 0
                     if num_per_channels is not None:
                         for i in T.Parallel(hidden_sf_aligned):
-                            if use_tma_aligned_col_major_sf:
-                                expanded_x_sf[i, pid_token] = 0
-                            else:
-                                expanded_x_sf[pid_token, i] = 0
+                            if i < hidden_sf:
+                                if use_tma_aligned_col_major_sf:
+                                    expanded_x_sf[i, pid_token] = 0
+                                else:
+                                    expanded_x_sf[pid_token, i] = 0
 
             if pid_token >= num_tokens:
                 T.thread_return()
@@ -80,13 +82,15 @@ def expand_to_fused_kernel(
                 T.assume(pos_local[k] < num_expanded_tokens)
                 if pos_local[k] >= 0:
                     for i in T.Parallel(hidden_aligned):
-                        expanded_x[pos_local[k], i] = x_fragment[i]
+                        if i < hidden:
+                            expanded_x[pos_local[k], i] = x_fragment[i]
                     if num_per_channels is not None:
                         for i in T.Parallel(hidden_sf_aligned):
-                            if use_tma_aligned_col_major_sf:
-                                expanded_x_sf[i, pos_local[k]] = x_sf_fragment[i]
-                            else:
-                                expanded_x_sf[pos_local[k], i] = x_sf_fragment[i]
+                            if i < hidden_sf:
+                                if use_tma_aligned_col_major_sf:
+                                    expanded_x_sf[i, pos_local[k]] = x_sf_fragment[i]
+                                else:
+                                    expanded_x_sf[pos_local[k], i] = x_sf_fragment[i]
 
     return expand_to_fused_kernel
 
diff --git a/tile_kernels/moe/get_fused_mapping_kernel.py b/tile_kernels/moe/get_fused_mapping_kernel.py
index 998189e..a4457cb 100644
--- a/tile_kernels/moe/get_fused_mapping_kernel.py
+++ b/tile_kernels/moe/get_fused_mapping_kernel.py
@@ -104,6 +104,7 @@ def get_fused_mapping_kernel(
             T.sync_grid()
 
             cumsum_shared = T.alloc_shared((num_threads,), T.int32)
+            cumsum_shared[thread_idx] = 0
             expert_num_elements = T.alloc_var(T.int32, init=0)
             expert_num_elements_aligned = T.alloc_var(T.int32, init=0)
             prefix_expert_num_elements = T.alloc_var(T.int32, init=0)
@@ -143,7 +144,9 @@ def get_fused_mapping_kernel(
             lane_mask_rev = ~lane_mask
             for i in T.serial(start + lane_idx, aligned_end, warp_size):
                 T.assume(0 <= i)
-                expert_idx = T.Select(i < numel, T.int32(topk_idx_1d[i]), -1)
+                expert_idx = T.alloc_var(T.int32, init=-1)
+                if i < numel:
+                    expert_idx = T.int32(topk_idx_1d[i])
                 mask = T.call_extern(T.uint32, '__match_any_sync', 0xFFFFFFFF, expert_idx)
                 count = T.popcount(mask & lane_mask)
 
@@ -201,7 +204,9 @@ def get_fused_mapping(
     should_sync = False
     if num_expanded_tokens == 0 and not force_no_sync:
         should_sync = True
-        num_expanded_tokens = (num_tokens * num_topk + (alignment - 1) * num_experts) // alignment * alignment
+        # Each expert range is aligned independently, so reserve the rounded-up
+        # upper bound before trimming with num_tokens_per_expert.
+        num_expanded_tokens = align(num_tokens * num_topk + (alignment - 1) * num_experts, alignment)
 
     # Allocate output
     num_sms = get_num_sms()