From 8166aa58aae2a405c4ed7ff26b4ab05c12af132d Mon Sep 17 00:00:00 2001
From: "PoYen, Chen" <PoYen.Chen@amd.com>
Date: Mon, 19 Aug 2024 08:14:45 +0000
Subject: [PATCH] Fix wrong uneven split checking logics

---
 example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
index c92e677e6c..ba826c8fb3 100644
--- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
+++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
@@ -114,8 +114,11 @@ template<>
 void fmha_fwd_splitkv_oneshot_<trait_{F_idx}>(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a)
 {{
     if constexpr({F_mode} == false) {{ // batch mode
+        // we don't check every seqlen_k values for kvcache
+        if (a.seqlen_k_ptr != nullptr) {{
+            kernel_runner<true>::run(s, a);
         // make sure F_bn0 is divisible by F_bk1
-        if (a.seqlen_k % (a.num_splits * {F_bn0}) == 0) {{
+        }} else if (a.seqlen_k % (a.num_splits * {F_bn0}) == 0) {{
             kernel_runner<false>::run(s, a);
         }} else {{
             kernel_runner<true>::run(s, a);