From 8b85fa6cf200b5f27ca8194df466e1adcef93014 Mon Sep 17 00:00:00 2001
From: Feng Shijie <Shijie.Feng@amd.com>
Date: Mon, 11 Aug 2025 06:03:06 +0000
Subject: [PATCH] update granularity control

---
 example/ck_tile/21_moe_flatmm/run_moe_flatmm_example.inc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/example/ck_tile/21_moe_flatmm/run_moe_flatmm_example.inc b/example/ck_tile/21_moe_flatmm/run_moe_flatmm_example.inc
index 0919334566..375da21232 100644
--- a/example/ck_tile/21_moe_flatmm/run_moe_flatmm_example.inc
+++ b/example/ck_tile/21_moe_flatmm/run_moe_flatmm_example.inc
@@ -234,7 +234,8 @@ int run_moe_gemm_example_with_layouts(int argc,
         static_cast<AccDataType*>(expert_weight_dev.GetDeviceBuffer());
 
     using MoeFlatmmArgs =
-        ck_tile::MoeFlatmmHostArgs<ck_tile::FlatmmScalePointer<1>, ck_tile::FlatmmScalePointer<1>>;
+        ck_tile::MoeFlatmmHostArgs<ck_tile::FlatmmScalePointer<ScaleGranularityM>,
+                                   ck_tile::FlatmmScalePointer<ScaleGranularityN>>;
 
     auto per_token_scale_dev_ptr = ck_tile::FlatmmScalePointer<ScaleGranularityM>{
         static_cast<float*>(per_token_scale_dev_buf.GetDeviceBuffer())};