From 55fc7330b83f6caa8cb83fc7b4b67609d9e99fef Mon Sep 17 00:00:00 2001
From: Qianfeng Zhang <Qianfeng.Zhang@amd.com>
Date: Fri, 5 Sep 2025 08:32:15 +0000
Subject: [PATCH] Move the loading of the first two v tiles into the k0-loop

---
 ...litkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp | 57 ++++++++++---------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp
index a3d979507b..cec253eb93 100644
--- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp
+++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp
@@ -398,38 +398,40 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
                     0); // prevent from messing up the order of global loads
             }
 
-            static_for<0, k0_loops, 1>{}([&](auto i_k0) {
-                block_sync_lds();
-                gemm_0(s_acc,
-                       get_slice_tile(
-                           q_tile, sequence<0, i_k0 * kK0>{}, sequence<kM0, (i_k0 + 1) * kK0>{}),
-                       k_lds_window);
-                block_sync_lds();
-
-                if constexpr(i_k0 < k0_loops - 1)
-                {
-                    store_tile(
-                        k_lds_window,
-                        tile_elementwise_in(k_element_func, k_tiles[number<(i_k0 + 1) % 2>{}]));
-                }
-
-                if constexpr(i_k0 < k0_loops - 2)
-                {
-                    k_tiles[number<i_k0 % 2>{}] = load_tile(k_dram_window);
-                    move_tile_window(k_dram_window, {0, kK0});
-                }
-            });
-
             using v_tile_type = decltype(load_tile(v_dram_window));
 
             statically_indexed_array<v_tile_type, 2> v_tiles;
 
-            static_for<0, 2, 1>{}(
-                [&, &i_page_block_v_ = i_page_block_v, &v_dram_window_ = v_dram_window](auto i_k1) {
-                    v_tiles[number<i_k1>{}] = load_tile(v_dram_window_);
+            static_for<0, k0_loops, 1>{}(
+                [&, &i_page_block_v_ = i_page_block_v, &v_dram_window_ = v_dram_window](auto i_k0) {
+                    block_sync_lds();
+                    gemm_0(s_acc,
+                           get_slice_tile(q_tile,
+                                          sequence<0, i_k0 * kK0>{},
+                                          sequence<kM0, (i_k0 + 1) * kK0>{}),
+                           k_lds_window);
+                    block_sync_lds();
 
-                    i_page_block_v_ = v_page_block_navigator.move_tile_window(
-                        i_page_block_v_, v_dram_window_, {0, kK1});
+                    if constexpr(i_k0 < k0_loops - 1)
+                    {
+                        store_tile(
+                            k_lds_window,
+                            tile_elementwise_in(k_element_func, k_tiles[number<(i_k0 + 1) % 2>{}]));
+                    }
+
+                    if constexpr(i_k0 < k0_loops - 2)
+                    {
+                        k_tiles[number<i_k0 % 2>{}] = load_tile(k_dram_window);
+                        move_tile_window(k_dram_window, {0, kK0});
+                    }
+                    else
+                    {
+
+                        v_tiles[number<i_k0 - (k0_loops - 2)>{}] = load_tile(v_dram_window_);
+
+                        i_page_block_v_ = v_page_block_navigator.move_tile_window(
+                            i_page_block_v_, v_dram_window_, {0, kK1});
+                    };
                 });
 
             // STAGE 2, scale_s, add bias, mask, softmax
@@ -567,6 +569,7 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
 
             // shuffle through LDS so that the tile layout is consistent with required by Gemm1
             store_tile(s_write_lds_window, s);
+            // at this point, all LDS operations are done including ds_read for k and ds_write for s
             block_sync_lds();
             auto s_new = load_tile(s_read_lds_window);