From 40a4d96cf5f07d6e649e660fade7fcfa2bf0f382 Mon Sep 17 00:00:00 2001
From: "PoYen, Chen" <PoYen.Chen@amd.com>
Date: Mon, 19 Aug 2024 10:16:23 +0000
Subject: [PATCH] Return earlier if split is empty

---
 .../pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp   | 6 ++++--
 1 file changed, 4 insertions(+), 2 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
index 2af54994a2..453c337143 100644
--- 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
@@ -223,13 +223,15 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
                 return seqlen_k_start_;
             }
         }();
-        const auto num_total_loop =
+        const index_t num_total_loop =
             integer_divide_ceil(seqlen_k_end - adjusted_seqlen_k_start, kN0);
 
         // check early exit if masked and no work to do.
         if constexpr(FmhaMask::IsMasking || kHasUnevenSplits)
         {
-            if(num_total_loop <= 0)
+            const index_t original_num_total_loop =
+                integer_divide_ceil(seqlen_k_end - seqlen_k_start, kN0);
+            if(original_num_total_loop <= 0)
             {
                 if constexpr(kStoreLSE)
                 {