diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp
index 66825edcf9..0633cc3eab 100644
--- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp
+++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp
@@ -356,7 +356,9 @@ int main(int argc, char* argv[])
                                KBatch,
                                a_element_op,
                                b_element_op,
-                               cde_element_op);
+                               cde_element_op,
+                               N,
+                               K);
 
     if(!device_op.IsSupportedArgument(argument))
     {
diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp
index c06e595c0f..b79889b991 100644
--- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp
+++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp
@@ -438,7 +438,9 @@ int main(int argc, char* argv[])
                                KBatch,
                                a_element_op,
                                b_element_op,
-                               cde_element_op);
+                               cde_element_op,
+                               N,
+                               K);
 
     if(!device_op.IsSupportedArgument(argument))
     {
diff --git a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8.cpp b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8.cpp
index 0d12441016..5abac7c8cf 100644
--- a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8.cpp
+++ b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8.cpp
@@ -365,7 +365,9 @@ int main(int argc, char* argv[])
                                KBatch,
                                a_element_op,
                                b_element_op,
-                               cde_element_op);
+                               cde_element_op,
+                               N,
+                               K);
 
     if(!device_op.IsSupportedArgument(argument))
     {
diff --git a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp
index c80b01d8c5..d3bf67ca86 100644
--- a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp
+++ b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp
@@ -405,7 +405,9 @@ int main(int argc, char* argv[])
                                KBatch,
                                a_element_op,
                                b_element_op,
-                               cde_element_op);
+                               cde_element_op,
+                               N,
+                               K);
 
     if(!device_op.IsSupportedArgument(argument))
     {
diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
index 950fe0236d..3d334c0fcf 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
@@ -391,7 +391,9 @@ struct DeviceMoeGemm : public DeviceGemmMultipleDSplitKBPreShuffle<ALayout,
                              index_t KBatch,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
-                             CElementwiseOperation c_element_op)
+                             CElementwiseOperation c_element_op,
+                             int,
+                             int)
     {
         return Argument{static_cast<const index_t*>(p_sorted_token_ids),
                         static_cast<const index_t*>(p_sorted_expert_ids),
@@ -432,7 +434,9 @@ struct DeviceMoeGemm : public DeviceGemmMultipleDSplitKBPreShuffle<ALayout,
                                                       index_t KBatch,
                                                       AElementwiseOperation a_element_op,
                                                       BElementwiseOperation b_element_op,
-                                                      CElementwiseOperation c_element_op) override
+                                                      CElementwiseOperation c_element_op,
+                                                      int,
+                                                      int) override
     {
         return std::make_unique<Argument>(nullptr,
                                           nullptr,