From 9d9c5a6c2464f2ad0abf9ee953de8d5bb84bb6bf Mon Sep 17 00:00:00 2001
From: "PoYen, Chen" <PoYen.Chen@amd.com>
Date: Thu, 8 Aug 2024 08:26:55 +0000
Subject: [PATCH] Fix compilation errors

---
 example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py | 2 +-
 example/ck_tile/01_fmha/fmha_fwd.cpp                    | 6 +++---
 example/ck_tile/01_fmha/fmha_fwd.hpp                    | 4 +++-
 3 files changed, 7 insertions(+), 5 deletions(-)

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 228bc5f3b7..f6a0e93262 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
@@ -220,7 +220,7 @@ float fmha_fwd_splitkv_(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a
     );
 }}
 
-float fmha_fwd_splitkv(fmha_fwd_traits t, fmha_fwd_splitkv_args a, const ck_tile::stream_config& s){{
+float fmha_fwd_splitkv(fmha_fwd_splitkv_traits t, fmha_fwd_splitkv_args a, const ck_tile::stream_config& s){{
     float r = -1;
 {F_dispatch}
     return r;
diff --git a/example/ck_tile/01_fmha/fmha_fwd.cpp b/example/ck_tile/01_fmha/fmha_fwd.cpp
index 7ab5c1d250..ba04c73342 100644
--- a/example/ck_tile/01_fmha/fmha_fwd.cpp
+++ b/example/ck_tile/01_fmha/fmha_fwd.cpp
@@ -989,12 +989,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
     const float fwd_ave_time = [&] {
 #if CK_TILE_FMHA_FWD_SPLITKV_API
-        if(1 < args.num_splits || args.block_table_ptr != nullptr)
+        if(1 < num_splits || 0 < page_block_size)
         {
-            return fmha_fwd_splitkv(traits, args, config);
+            return fmha_fwd_splitkv(fmha_traits, fmha_args, stream_config);
         }
 #endif
-        return fmha_fwd(traits, args, config);
+        return fmha_fwd(fmha_traits, fmha_args, stream_config);
     }();
 
     if(appendkv_ave_time < 0 || fwd_ave_time < 0)
diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp
index 36becf1b1a..54f426d14e 100644
--- a/example/ck_tile/01_fmha/fmha_fwd.hpp
+++ b/example/ck_tile/01_fmha/fmha_fwd.hpp
@@ -745,7 +745,9 @@ struct fmha_fwd_traits
     // TODO: padding check is inside this api
 };
 float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&);
-float fmha_fwd_splitkv(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&);
+
+using fmha_fwd_splitkv_traits = fmha_fwd_traits;
+float fmha_fwd_splitkv(fmha_fwd_splitkv_traits, fmha_fwd_splitkv_args, const ck_tile::stream_config&);
 
 struct fmha_fwd_appendkv_traits
 {