From 5ee7e2cf979c6c00d695a600d778c2f9633db093 Mon Sep 17 00:00:00 2001 From: Tianxing Wu Date: Thu, 11 Sep 2025 10:43:00 +0000 Subject: [PATCH] fixed var name --- .../ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp b/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp index c158be1ea1..2b6853553c 100644 --- a/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp +++ b/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp @@ -756,9 +756,9 @@ struct MoeFlatmmKernel // total number of tokens: sorted tokens + delimiter tokens + trailing padding tokens // we launch the grid based on the total number of tokens which needs to be static int partition_idx = blockIdx.x; - auto valid_token_cnt = kargs.p_max_token_id[0]; // sorted tokens + delimiter tokens - int total_valid_tile_cnt = TilePartitioner::GridSize(valid_token_cnt, kargs.N); - auto tilePartitioner = TilePartitioner{valid_token_cnt, kargs.N}; + auto max_token_id = kargs.p_max_token_id[0]; // sorted tokens + delimiter tokens + int total_valid_tile_cnt = TilePartitioner::GridSize(max_token_id, kargs.N); + auto tilePartitioner = TilePartitioner{max_token_id, kargs.N}; do { if (partition_idx >= total_valid_tile_cnt) { @@ -770,7 +770,7 @@ struct MoeFlatmmKernel this->operator()(kargs, block_offset_m, block_offset_n); partition_idx += gridDim.x; - } while(UsePersistentKernel && partition_idx < total_work_tile_cnt); + } while(UsePersistentKernel && partition_idx < total_valid_tile_cnt); } template