diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp
index 6f3a7e6357..6f0404a1ca 100644
--- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp
@@ -381,7 +381,6 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
         static_for<0, DS_READ_A_PREFETCH_STAGES, 1>{}([&](auto m0) {
             static_for<0, KRepeat, 1>{}([&](auto k0) {
                 static_for<0, KGroup, 1>{}([&](auto kg0) {
-                    // K = k0 × KGroup × k1 = k0 × kg0 × A_K1
                     a_thread_copy_.Run(a_block_desc_m0_m1_m2_k0_k1_k2,
                                        make_tuple(m0, I0, I0, Number<k0 * KGroup + kg0>{}, I0, I0),
                                        a_block_buf.At(I0),
diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle.hpp
index 6c4195e75d..f193b093d1 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle.hpp
@@ -860,35 +860,37 @@ struct DeviceGemmMultipleD_Xdl_CShuffle : public DeviceGemmMultipleD<ALayout,
 #endif
         if(desc.has_main_k_block_loop)
         {
-            GridwiseGemm::template Run<true>(p_a_grid,
-                                             p_b_grid,
-                                             p_ds_grid,
-                                             p_e_grid,
-                                             p_shared_block,
-                                             desc.a_element_op,
-                                             desc.b_element_op,
-                                             desc.cde_element_op,
-                                             desc.a_grid_desc_ak0_m_ak1,
-                                             desc.b_grid_desc_bk0_n_bk1,
-                                             desc.ds_grid_desc_mblock_mperblock_nblock_nperblock,
-                                             desc.e_grid_desc_mblock_mperblock_nblock_nperblock,
-                                             desc.block_2_etile_map);
+            GridwiseGemm::template Run<true, InMemoryDataOperationEnum::Set>(
+                p_a_grid,
+                p_b_grid,
+                p_ds_grid,
+                p_e_grid,
+                p_shared_block,
+                desc.a_element_op,
+                desc.b_element_op,
+                desc.cde_element_op,
+                desc.a_grid_desc_ak0_m_ak1,
+                desc.b_grid_desc_bk0_n_bk1,
+                desc.ds_grid_desc_mblock_mperblock_nblock_nperblock,
+                desc.e_grid_desc_mblock_mperblock_nblock_nperblock,
+                desc.block_2_etile_map);
         }
         else
         {
-            GridwiseGemm::template Run<false>(p_a_grid,
-                                              p_b_grid,
-                                              p_ds_grid,
-                                              p_e_grid,
-                                              p_shared_block,
-                                              desc.a_element_op,
-                                              desc.b_element_op,
-                                              desc.cde_element_op,
-                                              desc.a_grid_desc_ak0_m_ak1,
-                                              desc.b_grid_desc_bk0_n_bk1,
-                                              desc.ds_grid_desc_mblock_mperblock_nblock_nperblock,
-                                              desc.e_grid_desc_mblock_mperblock_nblock_nperblock,
-                                              desc.block_2_etile_map);
+            GridwiseGemm::template Run<false, InMemoryDataOperationEnum::Set>(
+                p_a_grid,
+                p_b_grid,
+                p_ds_grid,
+                p_e_grid,
+                p_shared_block,
+                desc.a_element_op,
+                desc.b_element_op,
+                desc.cde_element_op,
+                desc.a_grid_desc_ak0_m_ak1,
+                desc.b_grid_desc_bk0_n_bk1,
+                desc.ds_grid_desc_mblock_mperblock_nblock_nperblock,
+                desc.e_grid_desc_mblock_mperblock_nblock_nperblock,
+                desc.block_2_etile_map);
         }
     }
 };