diff --git a/example/65_gemm_multiply_multiply/moe_gemm1.cpp b/example/65_gemm_multiply_multiply/moe_gemm1.cpp index 0d98164bf1..e2abb4af6e 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm1.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1.cpp @@ -139,7 +139,7 @@ static constexpr ck::index_t BLOCKSIZE = 256; static constexpr ck::index_t NPerBlock = 128; static constexpr ck::index_t MNPerXDL = 32; static constexpr ck::index_t KPerBlock = 128 / sizeof(A0DataType); -static constexpr ck::index_t Nswizzle = true; +static constexpr ck::index_t Nswizzle = false; static constexpr ck::index_t AK1 = 16 / sizeof(A0DataType); static constexpr ck::index_t BK1 = 16 / sizeof(B0DataType); static constexpr ck::index_t EVec = 16 / sizeof(EDataType); 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 31a5c75083..61d7221097 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp @@ -1155,12 +1155,11 @@ struct GridwiseMoeGemm c_grid_desc_m_n, problem.MBlock, problem.NBlock); 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 expert_block_id = blockIdx.x / problem.NBlock; // const index_t b_block_id = blockIdx.x % problem.NBlock; - const index_t expert_id = __builtin_amdgcn_readfirstlane(p_sorted_expert_ids[expert_block_id]); const auto block_mn = [&]() -> std::pair { if constexpr (NSwizzle) { + const index_t expert_block_id = blockIdx.x / problem.NBlock; const index_t es = __builtin_amdgcn_readfirstlane(p_max_token_id[expert_block_id + 1]); const index_t expert_swizzle = es > 0 ? es : 1; //p_max_token_id[expert_id + 1]; const index_t expert_block_swizzle = expert_block_id / expert_swizzle; @@ -1174,6 +1173,7 @@ struct GridwiseMoeGemm }(); const index_t block_n_id = block_mn.first; const index_t block_m_id = block_mn.second; + const index_t expert_id = __builtin_amdgcn_readfirstlane(p_sorted_expert_ids[block_m_id]); // if (threadIdx.x==0) { // printf("bid %d, eid %d, es %d, esi %d, bsi %d, m %d, n %d\n", blockIdx.x, expert_id, expert_swizzle, expert_block_swizzle, b_block_id_swizzle, block_m_id, block_n_id); // }