Re-arrange V tile move_tile_window() statement

ROCm · Jan 2, 2025 · 25e1015 · 25e1015
1 parent b102083
commit 25e1015
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp
@@ -448,6 +448,7 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
                         });
                 }
             }
+
             __builtin_amdgcn_sched_barrier(0);
             // move K tile window
             i_page_block_k = k_page_block_navigator.move_tile_window(
@@ -466,6 +467,12 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
             tile_elementwise_inout(
                 [](auto& e0, auto e1, auto e2) { e0 = max(e1, e2); }, m, m_old, m_local); // m{j}
 
+            __builtin_amdgcn_sched_barrier(0);
+            // move V tile window
+            i_page_block_v = v_page_block_navigator.move_tile_window(
+                i_page_block_v, v_dram_block_window, {0, kN0});
+            __builtin_amdgcn_sched_barrier(0);
+
             auto p_compute = make_static_distributed_tensor<SMPLComputeDataType>(
                 s.get_tile_distribution()); // Pcompute{j}
 
@@ -565,12 +572,6 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
             const auto p =
                 cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, p_compute));
 
-            __builtin_amdgcn_sched_barrier(0);
-            // move V tile window
-            i_page_block_v = v_page_block_navigator.move_tile_window(
-                i_page_block_v, v_dram_block_window, {0, kN0});
-            __builtin_amdgcn_sched_barrier(0);
-
             // STAGE 3, KV gemm
             if constexpr(k1_loops > 1)
             {