diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp
index 01580b586d..b4eea2ad04 100644
--- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp
+++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.hpp
@@ -142,36 +142,90 @@ struct FlashAttentionFwd
     }
 };
 
-// TODO: change to only dec
-template <typename QDataType,
-          typename KDataType,
-          typename VDataType,
-          typename SaccDataType,
+// TODO: fwd_api.cpp
+template <typename SaccDataType_,
+          typename SMPLComputeDataType_,
+          typename PDataType_,
+          typename OaccDataType_,
+          index_t kBlockSize_,
+          index_t kHeadDim_,
+          index_t kM0PerBlock_,
+          index_t kN0PerBlock_,
+          index_t kK0PerBlock_,
+          index_t kN1PerBlock_,
+          index_t kK1PerBlock_>
+struct flash_attention_fwd_traits_
+{
+    using SaccDataType = ck_tile::remove_cvref_t<SaccDataType_>;
+    using SMPLComputeDataType = ck_tile::remove_cvref_t<SMPLComputeDataType_>;
+    using PDataType = ck_tile::remove_cvref_t<PDataType_>;
+    using OaccDataType = ck_tile::remove_cvref_t<OaccDataType_>;
+
+    static constexpr index_t kBlockSize  = kBlockSize_;
+    static constexpr index_t kHeadDim    = kHeadDim_;
+    static constexpr index_t kM0PerBlock = kM0PerBlock_;
+    static constexpr index_t kN0PerBlock = kN0PerBlock_;
+    static constexpr index_t kK0PerBlock = kK0PerBlock_;
+    static constexpr index_t kN1PerBlock = kN1PerBlock_;
+    static constexpr index_t kK1PerBlock = kK1PerBlock_;
+    
+    static constexpr ck_tile::index_t kWarpPerCu    = 8; // 2 warps per SIMD
+    static constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize;
+    static constexpr ck_tile::index_t kBlockPerCu   = kWarpPerCu / kWarpPerBlock;
+};
+
+// TODO: fwd_api.cpp, fwd_common.cpp
+template <typename SaccDataType,
           typename SMPLComputeDataType,
           typename PDataType,
           typename OaccDataType,
-          typename ODataType>
-float flash_attention_fwd(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& a, 
+          index_t kBlockSize,
+          index_t kHeadDim,
+          index_t kM0PerBlock,
+          index_t kN0PerBlock,
+          index_t kK0PerBlock,
+          index_t kN1PerBlock,
+          index_t kK1PerBlock>
+using traits_ = flash_attention_fwd_traits_<SaccDataType,
+                                            SMPLComputeDataType,
+                                            PDataType,
+                                            OaccDataType,
+                                            kBlockSize,
+                                            kHeadDim,
+                                            kM0PerBlock,
+                                            kN0PerBlock,
+                                            kK0PerBlock,
+                                            kN1PerBlock,
+                                            kK1PerBlock>;
+
+// Note: this internal API only declare, not define here, otherwise will block `make -j`
+template <typename QDataType,
+          typename KDataType,
+          typename VDataType,
+          typename ODataType,
+          typename Traits_>
+float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& a, 
+                          const ck_tile::stream_config& stream_config);
+
+// TODO: fwd_common.cpp
+template <typename QDataType,
+          typename KDataType,
+          typename VDataType,
+          typename ODataType,
+          typename Traits_>
+float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& a, 
                           const ck_tile::stream_config& stream_config) {
-    constexpr ck_tile::index_t kM0PerBlock = 128;
-    constexpr ck_tile::index_t kN0PerBlock = 128;
-    constexpr ck_tile::index_t kK0PerBlock = 32;
-    constexpr ck_tile::index_t kN1PerBlock = 128;
-    constexpr ck_tile::index_t kK1PerBlock = 32;
-
-    constexpr ck_tile::index_t kBlockSize = 256;
-    constexpr ck_tile::index_t kHeadDim   = 128;
-
-    ck_tile::index_t kGridSize = a.Batch * (a.M0 / kM0PerBlock) * (a.N1 / kN1PerBlock);
+    using SaccDataType        = typename Traits_::SaccDataType;                                                                           
+    using SMPLComputeDataType = typename Traits_::SMPLComputeDataType;                                                                           
+    using PDataType           = typename Traits_::PDataType;                                                                           
+    using OaccDataType        = typename Traits_::OaccDataType;                                                                           
+    
+    index_t kGridSize = a.Batch * (a.M0 / Traits_::kM0PerBlock) * (a.N1 / Traits_::kN1PerBlock);
 
     std::cout << "grid size " << kGridSize << std::endl;
 
-    constexpr ck_tile::index_t kWarpPerCu    = 8; // 2 warps per SIMD
-    constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize;
-    constexpr ck_tile::index_t kBlockPerCu   = kWarpPerCu / kWarpPerBlock;
-
     return ck_tile::launch_kernel(stream_config,
-        ck_tile::make_kernel<kBlockSize, kBlockPerCu>(
+        ck_tile::make_kernel<Traits_::kBlockSize, Traits_::kBlockPerCu>(
         ck_tile::FlashAttentionFwd<QDataType,
                                    KDataType,
                                    VDataType,
@@ -180,15 +234,15 @@ float flash_attention_fwd(const FlashAttnArgs<QDataType, KDataType, VDataType, O
                                    PDataType,
                                    OaccDataType,
                                    ODataType,
-                                   kBlockSize,
-                                   kHeadDim,
-                                   kM0PerBlock,
-                                   kN0PerBlock,
-                                   kK0PerBlock,
-                                   kN1PerBlock,
-                                   kK1PerBlock>{},
+                                   Traits_::kBlockSize,
+                                   Traits_::kHeadDim,
+                                   Traits_::kM0PerBlock,
+                                   Traits_::kN0PerBlock,
+                                   Traits_::kK0PerBlock,
+                                   Traits_::kN1PerBlock,
+                                   Traits_::kK1PerBlock>{},
         kGridSize,
-        kBlockSize,
+        Traits_::kBlockSize,
         0,
         a.q_ptr,
         a.k_ptr,
@@ -209,4 +263,44 @@ float flash_attention_fwd(const FlashAttnArgs<QDataType, KDataType, VDataType, O
         a.batchStrideO)); // BatchStrideO
 }
 
+// TODO: change to only declare
+// TODO: fwd_api.cpp
+template <typename QDataType,
+          typename KDataType,
+          typename VDataType,
+          typename SaccDataType,
+          typename SMPLComputeDataType,
+          typename PDataType,
+          typename OaccDataType,
+          typename ODataType>
+float flash_attention_fwd(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& a, 
+                          const ck_tile::stream_config& stream_config) {
+    constexpr ck_tile::index_t kM0PerBlock = 128;
+    constexpr ck_tile::index_t kN0PerBlock = 128;
+    constexpr ck_tile::index_t kK0PerBlock = 32;
+    constexpr ck_tile::index_t kN1PerBlock = 128;
+    constexpr ck_tile::index_t kK1PerBlock = 32;
+
+    constexpr ck_tile::index_t kBlockSize = 256;
+    constexpr ck_tile::index_t kHeadDim   = 128;
+
+    return flash_attention_fwd_<QDataType, 
+                                KDataType, 
+                                VDataType, 
+                                ODataType, 
+                                traits_<SaccDataType, 
+                                        SMPLComputeDataType, 
+                                        PDataType, 
+                                        OaccDataType,
+                                        kBlockSize, 
+                                        kHeadDim, 
+                                        kM0PerBlock, 
+                                        kN0PerBlock, 
+                                        kK0PerBlock, 
+                                        kN1PerBlock, 
+                                        kK1PerBlock>>
+            (a, stream_config);
+
+}
+
 } // namespace ck_tile