From aee38fdbf35c3bbf6f4badf0f05bf6146f2a0b27 Mon Sep 17 00:00:00 2001
From: yadaish <yadai@amd.com>
Date: Tue, 9 Dec 2025 00:30:52 +0000
Subject: [PATCH] update

---
 .../18_flatmm/mixed_prec/run_a16w4_moe_flatmm_example.inc      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/example/ck_tile/18_flatmm/mixed_prec/run_a16w4_moe_flatmm_example.inc b/example/ck_tile/18_flatmm/mixed_prec/run_a16w4_moe_flatmm_example.inc
index 38d685d4ab..116c55f333 100644
--- a/example/ck_tile/18_flatmm/mixed_prec/run_a16w4_moe_flatmm_example.inc
+++ b/example/ck_tile/18_flatmm/mixed_prec/run_a16w4_moe_flatmm_example.inc
@@ -72,7 +72,8 @@ int run_a16w4_moe_gemm_example_with_layouts(int argc,
     using CDataType   = PrecActType;
     using AccDataType = float;
 
-    using ScaleType = ck_tile::e8m0_t;
+    static constexpr bool IsInt4 = std::is_same_v<BDataType, ck_tile::pk_int4_t>;
+    using ScaleType = std::conditional_t<IsInt4, ck_tile::bfloat16_t, ck_tile::e8m0_t>;
 
     constexpr int ScaleGranularityN = 1;
     constexpr int ScaleGranularityK = 32;