diff --git a/include/ck_tile/core/tensor/tile_distribution.hpp b/include/ck_tile/core/tensor/tile_distribution.hpp
index 9392b9be48..8fe16a1693 100644
--- a/include/ck_tile/core/tensor/tile_distribution.hpp
+++ b/include/ck_tile/core/tensor/tile_distribution.hpp
@@ -114,22 +114,8 @@ struct tile_distribution
         }
         else if constexpr(NDimP == 2)
         {
-            return array<index_t, 2>{get_warp_id_to_sgpr(), get_lane_id()};
-        }
-    }
-
-    CK_TILE_HOST_DEVICE static auto _get_partition_index_v2()
-    {
-        // only support warp-tile and block-tile
-        static_assert(NDimP == 1 or NDimP == 2, "wrong!");
-
-        if constexpr(NDimP == 1)
-        {
-            return array<index_t, 1>{get_lane_id()};
-        }
-        else if constexpr(NDimP == 2)
-        {
-            return array<index_t, 2>{get_warp_id_to_vgpr(), get_lane_id()};
+            return array<index_t, 2>{get_warp_id(bool_constant<save_warp_id_in_sgpr>{}),
+                                     get_lane_id()};
         }
     }