diff --git a/example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp b/example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp
index 466420f066..f706cf19ff 100644
--- a/example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp
+++ b/example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp
@@ -32,7 +32,7 @@ float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_conf
         a.block_m,                                   // index_t unit_size;
         a.num_experts,                               // index_t num_experts;
         a.topk,                                      // index_t topk;
-        a.num_tokens * a.stride_token * o_data_bytes // index_t moe_buf_bytes;
+        static_cast<ck::long_index_t>(a.num_tokens) * a.stride_token * o_data_bytes // index_t moe_buf_bytes;
     };
 
     auto t1 = fused_moegemm_traits{t.prec_i,
diff --git a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp
index 3b1d91884f..56aadb9861 100644
--- a/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp
+++ b/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp
@@ -192,7 +192,7 @@ struct MoeSortingHostArgs
     index_t unit_size;      // this is the M_a of fused-moe kernel
     index_t num_experts;
     index_t topk;
-    index_t moe_buf_bytes;  // byte size of p_moe_buf
+    long_index_t moe_buf_bytes;  // byte size of p_moe_buf
 };
 
 template <typename Problem_>
@@ -219,7 +219,7 @@ struct MoeSortingKernel
         void* p_moe_buf;
         index_t tokens;
         index_t num_experts;
-        index_t moe_buf_bytes;
+        long_index_t moe_buf_bytes;
 
         index_t tokens_per_thread;
         index_t smem_rows;
@@ -426,7 +426,7 @@ struct MoeSortingKernel
         return row * total_col + col;
     }
 
-    CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes) const
+    CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, long_index_t buf_bytes) const
     {
         const index_t offset = (blockIdx.x - 1) * blockDim.x + threadIdx.x;
         if(offset < buf_bytes / 16)
@@ -1218,10 +1218,10 @@ CK_TILE_DEVICE void moe_sorting_wave_cumsum(data_t& thread_data)
 }
 
 template <index_t BLOCK_SIZE = 256>
-CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes, index_t gid)
+CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, long_index_t buf_bytes, index_t gid)
 {
     // const index_t offset = (blockIdx.x - 1) * BLOCK_SIZE + threadIdx.x;
-    index_t offset = gid * BLOCK_SIZE + threadIdx.x;
+    long_index_t offset = static_cast<long_index_t>(gid) * BLOCK_SIZE + threadIdx.x;
     if(offset < buf_bytes / 16)
     {
         buf[offset] = uint8x16_t{0};
@@ -1233,6 +1233,10 @@ CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes,
 // prefer to run mp kernel if is not oneshot
 CK_TILE_HOST bool moe_sorting_is_oneshot(int tokens_, int num_experts_)
 {
+    if(tokens_ >= 65536 * 2)
+    {
+        return 0;
+    }
     auto sub_token_          = moe_sorting_get_sub_token(tokens_, num_experts_);
     bool is_sub_token_onshot = tokens_ <= sub_token_;
     return is_sub_token_onshot;
@@ -1250,10 +1254,6 @@ CK_TILE_HOST index_t moe_sorting_mp_get_workspace_size(int tokens_, int num_expe
 CK_TILE_HOST index_t moe_sorting_get_workspace_size(int tokens_, int num_experts_)
 {
 #if 1
-    if(tokens_ >= 65536 * 2)
-    {
-        return 0;
-    }
     if(moe_sorting_is_oneshot(tokens_, num_experts_))
     {
         return 0;
@@ -1527,7 +1527,7 @@ struct MoeSortingMultiPhaseKernel_P2
         index_t num_experts;
         index_t mesh_stride; // mesh_stride for p_expert_mesh
         mdiv unit_size_mdiv;
-        index_t moe_buf_bytes;
+        long_index_t moe_buf_bytes;
     };
 
     CK_TILE_HOST static constexpr auto MakeKargs(const Hargs& h)