From a6241c62cc6be87a42e93dacfe54528db9a1ee09 Mon Sep 17 00:00:00 2001
From: Khushbu Agarwal <khuagarw@amd.com>
Date: Tue, 25 Nov 2025 13:10:28 -0800
Subject: [PATCH] [CK-Tile] fix block scale example for gfx1201 (#3283)

[ROCm/composable_kernel commit: 37ea1600888f515e5dfb7153b75b2f06474d880d]
---
 .../gemm_bquant_quantgrouped_preshuffleb.cpp   |  5 +++++
 ...uantgrouped_preshuffleb_preshufflequant.cpp |  5 +++++
 .../ck_tile/38_block_scale_gemm/gemm_utils.hpp | 18 ++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb.cpp
index 898316fa6b..8271559cdc 100644
--- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb.cpp
+++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb.cpp
@@ -3,8 +3,13 @@
 
 #include "run_gemm_quant_example.inc"
 
+#if CK_TILE_USE_WMMA
+template <typename T>
+using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill_Wmma<T>;
+#else
 template <typename T>
 using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill<T>;
+#endif
 
 void bquant_quantgrouped_preshuffleb_instance_factory(
     std::unordered_map<size_t, std::function<int(const ck_tile::ArgParser&)>>& lut)
diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant.cpp
index e26edb6501..f89591c791 100644
--- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant.cpp
+++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant.cpp
@@ -3,8 +3,13 @@
 
 #include "run_gemm_quant_example.inc"
 
+#if CK_TILE_USE_WMMA
+template <typename T>
+using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill_Wmma<T>;
+#else
 template <typename T>
 using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill<T>;
+#endif
 
 void bquant_quantgrouped_preshuffleb_preshufflequant_instance_factory(
     std::unordered_map<size_t, std::function<int(const ck_tile::ArgParser&)>>& lut)
diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
index 00118e4725..95b0a73ede 100644
--- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
+++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp
@@ -250,6 +250,24 @@ struct GemmConfigBQuantPrefill_Wmma : public GemmConfigBQuantPrefill<PrecType>
     static constexpr ck_tile::index_t K_Warp_Tile = 16;
 };
 
+template <typename PrecType>
+struct GemmConfigPreshuffleB_BQuant_Prefill_Wmma
+    : public GemmConfigPreshuffleB_BQuant_Prefill<PrecType>
+{
+    static constexpr ck_tile::index_t M_Warp_Tile = 16;
+    static constexpr ck_tile::index_t N_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile = 16;
+};
+
+template <typename PrecType>
+struct GemmConfigPreshuffleB_PreshuffleBQuant_Prefill_Wmma
+    : public GemmConfigPreshuffleB_PreshuffleBQuant_Prefill<PrecType>
+{
+    static constexpr ck_tile::index_t M_Warp_Tile = 16;
+    static constexpr ck_tile::index_t N_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile = 16;
+};
+
 template <typename ADataType_,
           typename BDataType_ = ADataType_,
           typename CDataType_ = ADataType_,