From 8d396d29f02be059c90d4cf3a8a11d7e238c95ec Mon Sep 17 00:00:00 2001
From: Amir Ghamarian <amir.ghamarian@amd.com>
Date: Sat, 28 Mar 2026 10:47:45 +0000
Subject: [PATCH] Add async prefetch overlap to single-warp-group pipeline

Move next iteration's K/V global loads (K_mem_load, V_mem_load) to
immediately after the barrier, before PV GEMM and K LDS read. This
overlaps the async global->LDS copies with the current iteration's
GEMM compute. Also remove redundant barriers between PV and QK phases
since K/V use separate LDS regions (no read/write conflicts).

Benchmark improvement (64-seq decode, d64 GQA-8):
  Phase 1: 0.03564ms -> Phase 2: 0.03406ms (~4.6% faster)
  Total vs original baseline: 0.06177ms -> 0.03406ms (1.81x speedup)

Made-with: Cursor
---
 .../pipeline/unified_attention_pipeline.hpp   | 40 +++++++++----------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
index 52f309a8a0..d91a122634 100644
--- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
+++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
@@ -962,7 +962,7 @@ struct UnifiedAttentionPipeline
         {
             if constexpr(NumWarpGroups == 1)
             {
-                // --- Single warp group: serial pipeline ---
+                // --- Single warp group: serial pipeline with async prefetch ---
                 // After pre-stage:
                 //   sp(0) has QK for block 0 (alu0 + alu_D_upd done, alu1 NOT done)
                 //   V0 loading to LDS (V buf 0)
@@ -973,17 +973,16 @@ struct UnifiedAttentionPipeline
                 s_waitcnt_vmcnt<0>();
                 __builtin_amdgcn_s_barrier();
 
-                V_lds_load(number<0>{}); // V0 from LDS
+                V_mem_load(number<1>{}); // prefetch V1 -> buf 1 (overlaps with compute)
+
+                V_lds_load(number<0>{}); // V0 from LDS -> kv_tile.v_tile
                 s_waitcnt_lgkmcnt<0>();
                 fmha_alu1(number<0>{}); // finalize sp(0) -> P(0)
                 gemm(number<0>{}, /*gemm_idx=*/number<1>{}); // PV: P(0)*V0
 
-                __builtin_amdgcn_s_barrier();
-                K_lds_load(number<1>{}); // K1 from LDS
+                K_lds_load(number<1>{}); // K1 from LDS -> kv_tile.k_tile
                 s_waitcnt_lgkmcnt<0>();
 
-                V_mem_load(number<1>{}); // start V1 -> LDS buf 1
-
                 gemm(number<1>{}, /*gemm_idx=*/number<0>{}); // QK: Q*K1 -> sp(1)
                 fmha_mask(number<1>{});
                 fmha_alu0(number<1>{});
@@ -997,19 +996,20 @@ struct UnifiedAttentionPipeline
                     s_waitcnt_vmcnt<0>();
                     __builtin_amdgcn_s_barrier();
 
-                    V_lds_load(number<1>{}); // V from buf 1 -> kv_tile.v_tile
+                    // Prefetch next iteration's K/V (overlaps with all compute below)
+                    // K/V use separate LDS regions so no conflict with current reads
+                    if(i_total_loops + 1 < num_total_loop)
+                        K_mem_load(number<1>{}); // next K -> K buf 1
+                    V_mem_load(number<0>{}); // next V -> V buf 0
+
+                    V_lds_load(number<1>{}); // V from V buf 1 -> kv_tile.v_tile
                     s_waitcnt_lgkmcnt<0>();
                     fmha_alu1(number<1>{}); // finalize sp(1) -> P(1)
                     gemm(number<1>{}, /*gemm_idx=*/number<1>{}); // PV: P(1)*V
 
-                    __builtin_amdgcn_s_barrier();
-                    K_lds_load(number<0>{}); // K from buf 0 -> kv_tile.k_tile
+                    K_lds_load(number<0>{}); // K from K buf 0 -> kv_tile.k_tile
                     s_waitcnt_lgkmcnt<0>();
 
-                    if(i_total_loops + 1 < num_total_loop)
-                        K_mem_load(number<1>{}); // next K -> buf 1
-                    V_mem_load(number<0>{}); // next V -> buf 0
-
                     gemm(number<0>{}, /*gemm_idx=*/number<0>{}); // QK -> sp(0)
                     fmha_mask(number<0>{});
                     fmha_alu0(number<0>{});
@@ -1023,19 +1023,19 @@ struct UnifiedAttentionPipeline
                     s_waitcnt_vmcnt<0>();
                     __builtin_amdgcn_s_barrier();
 
-                    V_lds_load(number<0>{}); // V from buf 0 -> kv_tile.v_tile
+                    // Prefetch next iteration's K/V
+                    if(i_total_loops + 1 < num_total_loop)
+                        K_mem_load(number<0>{}); // next K -> K buf 0
+                    V_mem_load(number<1>{}); // next V -> V buf 1
+
+                    V_lds_load(number<0>{}); // V from V buf 0 -> kv_tile.v_tile
                     s_waitcnt_lgkmcnt<0>();
                     fmha_alu1(number<0>{}); // finalize sp(0) -> P(0)
                     gemm(number<0>{}, /*gemm_idx=*/number<1>{}); // PV: P(0)*V
 
-                    __builtin_amdgcn_s_barrier();
-                    K_lds_load(number<1>{}); // K from buf 1 -> kv_tile.k_tile
+                    K_lds_load(number<1>{}); // K from K buf 1 -> kv_tile.k_tile
                     s_waitcnt_lgkmcnt<0>();
 
-                    if(i_total_loops + 1 < num_total_loop)
-                        K_mem_load(number<0>{}); // next K -> buf 0
-                    V_mem_load(number<1>{}); // next V -> buf 1
-
                     gemm(number<1>{}, /*gemm_idx=*/number<0>{}); // QK -> sp(1)
                     fmha_mask(number<1>{});
                     fmha_alu0(number<1>{});