From 6c8cfa58c490d48e51cd780f604347a6e8b71ff3 Mon Sep 17 00:00:00 2001
From: yadaish <yadai@amd.com>
Date: Tue, 25 Nov 2025 16:21:37 +0000
Subject: [PATCH] update

---
 include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp | 2 +-
 1 file changed, 1 insertion(+), 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 15c4c21c86..946b81c146 100644
--- a/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
+++ b/include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
@@ -1264,7 +1264,7 @@ struct MoeFlatmmKernel
                         scatter_token_id =
                             scatter_token_id * kargs.TopK + (fused_token >> token_id_offset);
                     c_scatter_offsets[mIter][m0] = scatter_token_id * kargs.stride_C;
-                    c_scatter_valids[mIter][m0]  = (scatter_token_id < kargs.NumTokens);
+                    c_scatter_valids[mIter][m0]  = (scatter_token_id < (kargs.NumTokens * (IsInputGemm? kargs.TopK : 1)));
                 });
             });