diff --git a/example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt b/example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
index 23fc7484dd..5af6c5af39 100644
--- a/example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
+++ b/example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
@@ -7,6 +7,8 @@ add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL flash_attention_fwd.cpp)
 target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
 set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
 
+# list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
+
 # NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
 list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
 
@@ -16,6 +18,12 @@ if(ENABLE_TOY_FA_FWD_OPT)
     target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE TOY_FA_FWD_OPT)
 endif()
 
+option(ENABLE_TOY_FA_FWD_QK_SWIZZLE "Enable toy FA fwd QK swizzle" OFF)
+if(ENABLE_TOY_FA_FWD_QK_SWIZZLE)
+    message("Compiling with toy FA fwd QK swizzle")
+    target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE TOY_FA_FWD_QK_SWIZZLE)
+endif()
+
 target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS})
 
 # TODO: we have to turn off this global prop, otherwise the progress bar generated
diff --git a/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp b/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
index 8994689841..9d794cd942 100644
--- a/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
+++ b/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
@@ -23,7 +23,12 @@ struct BlockGemmARegBSmemCRegV1DefaultPolicy
         }
         else if constexpr(kM0 == 128)
         {
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
             return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
+#else
+#pragma message("Enable toy FA fwd QK swizzle")
+            return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
+#endif
         }
         else
         {
diff --git a/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp b/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
index e3f3fd0cd6..5f6c56d0e1 100644
--- a/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
+++ b/example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
@@ -23,7 +23,11 @@ struct BlockGemmARegBSmemCRegV1K8Policy
         }
         else if constexpr(kM0 == 128)
         {
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
             return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
+#else
+            return make_tuple(WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution{}, 4, 1);
+#endif
         }
         else
         {
diff --git a/example/ck_tile/99_toy_example/03_flash_attention_fwd/flash_attention_fwd_impl.hpp b/example/ck_tile/99_toy_example/03_flash_attention_fwd/flash_attention_fwd_impl.hpp
index 5a9041973a..90779420c1 100644
--- a/example/ck_tile/99_toy_example/03_flash_attention_fwd/flash_attention_fwd_impl.hpp
+++ b/example/ck_tile/99_toy_example/03_flash_attention_fwd/flash_attention_fwd_impl.hpp
@@ -63,7 +63,11 @@ struct FlashAttentionFwdImpl
     {
         constexpr index_t kNPerBlock = kN1PerBlock;
         constexpr index_t kKPerBlock = kK1PerBlock;
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
         constexpr index_t kKPack     = 4;
+#else
+        constexpr index_t kKPack     = 8;
+#endif
 
         constexpr auto dataTypeSize = sizeof(VDataType);
         constexpr auto NLdsLayer =
diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
index 4f81f8bb70..4a114d4af3 100644
--- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
+++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
@@ -6,6 +6,7 @@ if(FLASH_ATTENTION_FWD_ENABLE_APIS STREQUAL "all")
 endif()
 
 option(TOY_FA_FWD_OPT "Enable toy flash attention forward optimization" ON)
+option(TOY_FA_FWD_QK_SWIZZLE "Enable toy flash attention forward QK swizzle" OFF)
 
 execute_process(
   COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py
diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
index 8994689841..9d794cd942 100644
--- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
+++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
@@ -23,7 +23,12 @@ struct BlockGemmARegBSmemCRegV1DefaultPolicy
         }
         else if constexpr(kM0 == 128)
         {
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
             return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
+#else
+#pragma message("Enable toy FA fwd QK swizzle")
+            return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
+#endif
         }
         else
         {
diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
index e3f3fd0cd6..5f6c56d0e1 100644
--- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
+++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
@@ -23,7 +23,11 @@ struct BlockGemmARegBSmemCRegV1K8Policy
         }
         else if constexpr(kM0 == 128)
         {
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
             return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
+#else
+            return make_tuple(WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution{}, 4, 1);
+#endif
         }
         else
         {
diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp
index 5a9041973a..90779420c1 100644
--- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp
+++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp
@@ -63,7 +63,11 @@ struct FlashAttentionFwdImpl
     {
         constexpr index_t kNPerBlock = kN1PerBlock;
         constexpr index_t kKPerBlock = kK1PerBlock;
+#if !defined(TOY_FA_FWD_QK_SWIZZLE)
         constexpr index_t kKPack     = 4;
+#else
+        constexpr index_t kKPack     = 8;
+#endif
 
         constexpr auto dataTypeSize = sizeof(VDataType);
         constexpr auto NLdsLayer =