From bd6897d4323dbbd3b13130111063c93d7d6d640d Mon Sep 17 00:00:00 2001
From: yadaish <yadai@amd.com>
Date: Thu, 4 Dec 2025 18:28:17 +0000
Subject: [PATCH] update

---
 include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp b/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
index 8ccd2c1bcb..e6f3364bcb 100644
--- a/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
+++ b/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
@@ -639,7 +639,8 @@ struct MoeFlatmmKernel
         index_t FlatScaleN = kargs.N / N_Pack / BlockGemmShape::WarpTile::at(I1);
 
         // using ScaleType = std::conditional_t<MXFP4_Pipeline, e8m0_t, float>;
-        using ScaleType = std::conditional_t<MXFP4_Pipeline, bfloat16_t, float>;
+        static constexpr bool IsInt4 = std::is_same_v<BDataType, pk_int4_t>;
+        using ScaleType = std::conditional_t<MXFP4_Pipeline, std::conditional_t<IsInt4, bfloat16, e8m0_t>, float>;
 
         const auto scale_b_flat_view = make_naive_tensor_view<address_space_enum::global>(
             reinterpret_cast<const ScaleType*>(scale_n.ptr) + expert_id * kargs.N * scale_k,