From 70aef2ed33d31071cbfc9cdea7c636f872bb6e44 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Thu, 20 Feb 2025 13:54:44 +0000 Subject: [PATCH] hot fix expert over flow --- include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp index 1e0fd23525..2ceb6bde9c 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp @@ -1174,7 +1174,9 @@ struct GridwiseMoeGemm const index_t max_token_id = __builtin_amdgcn_readfirstlane(p_max_token_id[0]); // constexpr int expert_tile_cnt[8] = {2, 1, 1, 2, 2, 2, 1, 2}; // const index_t b_block_id = blockIdx.x % problem.NBlock; - const index_t expert_block_id = blockIdx.x / problem.NBlock; + const index_t expert_block_id = NSwizzle ? blockIdx.x / problem.NBlock : blockIdx.x; + if (expert_block_id * MPerBlock >= max_token_id) + return; const index_t expert_id = NSwizzle ? __builtin_amdgcn_readfirstlane(p_sorted_expert_ids[blockIdx.x / problem.NBlock]) : __builtin_amdgcn_readfirstlane(p_sorted_expert_ids[blockIdx.y]); const auto block_mn = [&]() -> std::pair { @@ -1219,7 +1221,7 @@ struct GridwiseMoeGemm constexpr auto AMRepeats = MPerBlock / AMThreads; const index_t token_pos = block_m_id * MPerBlock + threadIdx.x / AKThreads * AMRepeats; - if(token_pos >= max_token_id || expert_block_id * MPerBlock >= max_token_id || token0 >= problem.NumTokens) + if(token_pos >= max_token_id || token0 >= problem.NumTokens) return; StaticallyIndexedArray gather_offsets; //= p_sorted_token_ids[token_pos]; static_for<0, AMRepeats, 1>{}([&](auto m0) {