diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py
index d184ef6..85de881 100644
--- a/deepinfer/neuron_ops/static_moe.py
+++ b/deepinfer/neuron_ops/static_moe.py
@@ -251,7 +251,7 @@ def _nki_moe_forward(
         x = packed[:, : x.shape[1]].to(torch.bfloat16)
         idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32)
         w = packed[:, x.shape[1] + topk :]
-    regions.split(coarse=True)
+    regions.split()
     local_idx = idx - moe.neuron_rank_tensor() * num_local
     overflow = None
     if capacity and capacity < x.shape[0]:
@@ -265,10 +265,10 @@ def _nki_moe_forward(
                 gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value
             )
             acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float()
-    regions.split(coarse=True)
+    regions.split()
     hier = getattr(moe, "_neuron_hier_ar", None)
     out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
-    regions.split(coarse=True)
+    regions.split()
     if moe.has_shared:
         shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
         shared = swiglu_forward_and_per_token_cast(