diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp index fbe82a5e86..db2359f298 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp @@ -178,7 +178,7 @@ static constexpr ck::index_t MPerBlock = 64; using DeviceOpInstance = ck::tensor 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 1, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 1, + S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 4, 2, S<1, 32, 1, 8>, S<2, 1, 1, 1>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ActOP, Nswizzle, true, MulRoutedWeight, int32_t, A0DataType>; #endif @@ -386,8 +386,8 @@ int main(int argc, char* argv[]) DeviceMem d2_device_buf(sizeof(D2DataType) * d2_e_n.mDesc.GetElementSpaceSize()); DeviceMem e_device_buf(sizeof(EDataType) * e_t_n_device_result.mDesc.GetElementSpaceSize()); // a0_t_k.savetxt("a.txt"); - expert_ids.savetxt("expert_ids.txt", "int"); - sorted_token_ids.savetxt("sorted_token_ids.txt", "int"); + // expert_ids.savetxt("expert_ids.txt", "int"); + // sorted_token_ids.savetxt("sorted_token_ids.txt", "int"); // d2_e_n.savetxt("d2_e_n.txt", "int"); sorted_token_ids_dev.ToDevice(sorted_token_ids.mData.data()); expert_ids_dev.ToDevice(expert_ids.mData.data()); @@ -517,6 +517,7 @@ int main(int argc, char* argv[]) PassThrough{}); ref_invoker.Run(ref_argument); + bool host_all_zero = true; for(int m = 0; m < valid_size; ++m) { @@ -532,6 +533,7 @@ int main(int argc, char* argv[]) { e_t_n_host_result(t, topk_id, n) = ck::type_convert(c_t_k_n(t, topk_id, n)); + host_all_zero = host_all_zero && !e_t_n_host_result(t, topk_id, n); } } @@ -566,6 +568,9 @@ int main(int argc, char* argv[]) : 1; if(status == 0) { + if (host_all_zero){ + printf("WARNING: All results are zero. "); + } printf("Validation Pass.\n"); } return status; diff --git a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp index 7e8697492f..0a88c0a27a 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp @@ -172,7 +172,7 @@ static constexpr ck::index_t MPerBlock = 64; using DeviceOpInstance = ck::tensor 16, 16, 16, 16, 4, 2, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 2, 2, S<1, 32, 1, 8>, S<2, 1, 1, 1>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, 0, false, false, MulRoutedWeight, int32_t, A0DataType>; @@ -184,24 +184,15 @@ int main(int argc, char* argv[]) bool do_verification = true; int init_method = 1; bool time_kernel = true; - - // tokens = 1 - // topk = 1 - // experts = 8 - // per expert: - - constexpr ck::index_t valid_tile_num = - 26; // 13 for 128; 52 for 32; 4096 for ds // > token * topk / MPerBlock - constexpr ck::index_t sorted_tile_num = valid_tile_num + 3; - ck::index_t sorted_size = sorted_tile_num * MPerBlock; - ck::index_t valid_size = valid_tile_num * MPerBlock; #if 1 // GEMM shape ck::index_t N = 6144; ck::index_t K = 4096; ck::index_t experts = 8; - ck::index_t tokens = 832; ck::index_t topk = 2; + ck::index_t valid_tile_num = 128; + ck::index_t sorted_tile_num = valid_tile_num + 3; + ck::index_t tokens = 4096; #else // deepseek ck::index_t N = 2048; @@ -231,6 +222,18 @@ int main(int argc, char* argv[]) K = std::stoi(argv[5]); tokens = std::stoi(argv[6]); } + else if(argc == 9) + { + + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + N = std::stoi(argv[4]); + K = std::stoi(argv[5]); + tokens = std::stoi(argv[6]); + sorted_tile_num = std::stoi(argv[7]); + valid_tile_num = std::stoi(argv[8]); + } else { printf("arg1: verification (0=no, 1=yes)\n"); @@ -240,6 +243,15 @@ int main(int argc, char* argv[]) exit(0); } + valid_tile_num = tokens * topk / MPerBlock; + sorted_tile_num = valid_tile_num + 3; + ck::index_t sorted_size = sorted_tile_num * MPerBlock; + ck::index_t valid_size = valid_tile_num * MPerBlock; + if(tokens * topk > valid_size) + { + printf("err config, tokens * topk > valid_size\n"); + exit(-1); + } ck::index_t StrideA = K; ck::index_t StrideB = K; ck::index_t StrideE = N; @@ -270,12 +282,8 @@ int main(int argc, char* argv[]) { expert_ids.mData[i] = i / ck::math::integer_divide_ceil(valid_tile_num, experts); } - if(tokens * topk > valid_size) - { - printf("err config, tokens * topk > valid_size\n"); - exit(-1); - } - int token_per_tile = tokens * topk / valid_tile_num; + + int token_per_tile = (tokens * topk + valid_tile_num - 1) / valid_tile_num; int tokenid = 0; for(int i = 0; i < sorted_size; i++) @@ -448,7 +456,7 @@ int main(int argc, char* argv[]) float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel}); std::size_t flop = std::size_t(2) * tokens * topk * N * K; - std::size_t num_btype = sizeof(A0DataType) * tokens * K * topk + + std::size_t num_btype = sizeof(A0DataType) * valid_tile_num * K + sizeof(B0DataType) * K * N * experts + sizeof(EDataType) * tokens * N; @@ -520,12 +528,14 @@ int main(int argc, char* argv[]) cde_element_op); ref_invoker.Run(ref_argument); + bool host_all_zero = true; for(int t = 0; t < tokens; ++t) { for(int n = 0; n < N; ++n) { e_t_n_host_result(t, n) = ck::type_convert(c_t_n(t, n)); + host_all_zero = host_all_zero && !e_t_n_host_result(t, n); } } @@ -561,6 +571,9 @@ int main(int argc, char* argv[]) : 1; if(status == 0) { + if (host_all_zero){ + printf("WARNING: All results are zero. "); + } printf("Validation Pass.\n"); } return status; diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp index a221d67f61..ef6e3d07fc 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp @@ -217,7 +217,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v1< static_for<0, num_buffer_load_inst_a, 1>{}([&](auto i) { ignore = i; __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - // __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read }); @@ -443,6 +442,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v1< // Local prefetch A1 __builtin_amdgcn_s_waitcnt(3952); + block_sync_lds(); static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { static_for<0, KGroup, 1>{}([&](auto kg0) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp index bf41ccbc19..9d2d05013f 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp @@ -201,12 +201,10 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< ? HotLoopInstList::A_LDS_Read_Inst_Num : HotLoopInstList::A_LDS_Read_Inst_Num / 2; - constexpr auto num_ds_write_inst_a = HotLoopInstList::A_LDS_Write_Inst_Num; constexpr auto num_buffer_load_inst_a = HotLoopInstList::A_Buffer_Load_Inst_Num; constexpr auto num_buffer_load_inst_b = HotLoopInstList::B_Buffer_Load_Inst_Num * 2; - static_assert(num_buffer_load_inst_a == num_ds_write_inst_a); constexpr auto num_mfma_inst = HotLoopInstList::C_MFMA_Inst_Num * 2; constexpr auto mfma_cycle = HotLoopInstList::C_MFMA_Inst_Cycle; @@ -257,7 +255,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< num_mfma_perstage / buffer_load_perstage_more ? num_mfma_perstage / buffer_load_perstage_more : 1; constexpr auto buffer_load_issue_point_interval_less = num_mfma_perstage / buffer_load_perstage_less ? num_mfma_perstage / buffer_load_perstage_less : 1; - constexpr auto ds_write_issue_point = 0; constexpr auto buffer_load_issue_point_a = num_mfma_perstage >= 3 ? 1 : 0; // B global read @@ -297,15 +294,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read static_for<0, num_mfma_perstage, 1>{}([&](auto imfma) { __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - if constexpr((((i + buffer_load_b_stages) < buffer_load_stages_more) && - (imfma % buffer_load_issue_point_interval_more == - ds_write_issue_point)) || - (((i + buffer_load_b_stages) >= buffer_load_stages_more) && - (imfma % buffer_load_issue_point_interval_less == - ds_write_issue_point))) - { - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write - } if constexpr((((i + buffer_load_b_stages) < buffer_load_stages_more) && (imfma % buffer_load_issue_point_interval_more == buffer_load_issue_point_a)) || @@ -426,7 +414,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< StaticallyIndexedArray{}> b_thread_bufs_up; constexpr auto b_block_origin_idx = make_tuple(I0, I0, I0, I0); - auto a_scale_thread_buf = make_static_buffer( + auto a_scale_thread_buf = make_static_buffer( a_scale_thread_desc.GetElementSpaceSize()); auto b_scale_thread_buf = make_static_buffer( b_scale_thread_desc.GetElementSpaceSize()); @@ -455,7 +443,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); b_blockwise_copy_up.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(I0)); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); __builtin_amdgcn_sched_barrier(0); @@ -499,12 +487,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< c_scale_thread_buf_up(m0) = a_scale_thread_bufs[I0][m0] * b_scale_thread_bufs_up[I0][I0]; }); - // Local prefill A1 - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(I0)); - // Global prefetch A2, AScale2 BScale2 - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); - a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); a_scale_thread_copy.Run(a_scale_grid_desc, a_scale_grid_buf, @@ -560,6 +543,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< c_thread_buf_per_scale_up; // Local prefetch A1 + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, 2, 1>{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { @@ -596,9 +580,9 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< b_block_origin_idx, b_thread_bufs_up(local_read_buf)); b_blockwise_copy_up.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); - - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(local_read_buf)); - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + + block_sync_lds(); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(local_read_buf)); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); a_scale_thread_copy.Run(a_scale_grid_desc, @@ -719,6 +703,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< if constexpr(m0.value == (MRepeat - 2)) { + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, KRepeat, 1>{}([&](auto k0) { @@ -828,7 +813,8 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< b_block_desc_n0_n1_k0_k1, b_block_origin_idx, b_thread_bufs_up(I1)); - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(I1)); + block_sync_lds(); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(I1)); static_for<0, MRepeat, 1>{}([&](auto m0) { vector_type c_scale_thread_vec; @@ -904,6 +890,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3< if constexpr(m0.value == (MRepeat - 2)) { + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, KRepeat, 1>{}([&](auto k0) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp index 9c3d6e4009..8464161e0e 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp @@ -209,7 +209,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< static_for<0, num_buffer_load_inst_a, 1>{}([&](auto i) { ignore = i; __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read }); @@ -301,7 +300,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< c_scale_thread_desc.GetElementSpaceSize()); // Global prefetch A1 B1 - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, I0); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf); b_blockwise_copy.Run(b_grid_desc, b_grid_buf, b_block_desc_n0_n1_k0_k1, @@ -359,12 +358,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< }); }); - // Local prefill A1 - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf, I0); - - // Global prefetch A2 - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, I0); - a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); a_scale_thread_copy.Run(a_scale_grid_desc, a_scale_grid_buf, @@ -399,6 +392,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< c_thread_buf_per_scale; // Local prefetch A1 + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { @@ -434,9 +428,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); block_sync_lds(); - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf, mfma_reg_buf); - - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, local_read_buf); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); static_for<0, MRepeat, 1>{}([&](auto m0) { @@ -512,7 +504,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< }); }); }); - + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, MRepeat, 1>{}([&](auto m0) { @@ -592,7 +584,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< b_block_origin_idx, b_thread_bufs(I1)); block_sync_lds(); - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf); static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, NRepeat, 1>{}([&](auto n0) { @@ -674,7 +666,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1< }); }); }); - + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, MRepeat, 1>{}([&](auto m0) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp index 406e8737e7..9388250648 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp @@ -201,12 +201,10 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< ? HotLoopInstList::A_LDS_Read_Inst_Num : HotLoopInstList::A_LDS_Read_Inst_Num / 2; - constexpr auto num_ds_write_inst_a = HotLoopInstList::A_LDS_Write_Inst_Num; constexpr auto num_buffer_load_inst_a = HotLoopInstList::A_Buffer_Load_Inst_Num; constexpr auto num_buffer_load_inst_b = HotLoopInstList::B_Buffer_Load_Inst_Num; - static_assert(num_buffer_load_inst_a == num_ds_write_inst_a); constexpr auto num_mfma_inst = HotLoopInstList::C_MFMA_Inst_Num; constexpr auto mfma_cycle = HotLoopInstList::C_MFMA_Inst_Cycle; @@ -257,7 +255,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< num_mfma_perstage / buffer_load_perstage_more ? num_mfma_perstage / buffer_load_perstage_more : 1; constexpr auto buffer_load_issue_point_interval_less = num_mfma_perstage / buffer_load_perstage_less ? num_mfma_perstage / buffer_load_perstage_less : 1; - constexpr auto ds_write_issue_point = 0; constexpr auto buffer_load_issue_point_a = num_mfma_perstage >= 3 ? 1 : 0; // B global read @@ -297,15 +294,6 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read static_for<0, num_mfma_perstage, 1>{}([&](auto imfma) { __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - if constexpr((((i + buffer_load_b_stages) < buffer_load_stages_more) && - (imfma % buffer_load_issue_point_interval_more == - ds_write_issue_point)) || - (((i + buffer_load_b_stages) >= buffer_load_stages_more) && - (imfma % buffer_load_issue_point_interval_less == - ds_write_issue_point))) - { - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write - } if constexpr((((i + buffer_load_b_stages) < buffer_load_stages_more) && (imfma % buffer_load_issue_point_interval_more == buffer_load_issue_point_a)) || @@ -420,7 +408,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< StaticallyIndexedArray{}> b_thread_bufs; constexpr auto b_block_origin_idx = make_tuple(I0, I0, I0, I0); - auto a_scale_thread_buf = make_static_buffer( + auto a_scale_thread_buf = make_static_buffer( a_scale_thread_desc.GetElementSpaceSize()); auto b_scale_thread_buf = make_static_buffer( b_scale_thread_desc.GetElementSpaceSize()); @@ -439,7 +427,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< b_thread_bufs(I0)); b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(I0)); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); __builtin_amdgcn_sched_barrier(0); @@ -472,12 +460,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< c_scale_thread_buf(m0) = a_scale_thread_bufs[I0][m0] * b_scale_thread_bufs[I0][I0]; }); - // Local prefill A1 - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(I0)); - // Global prefetch A2, AScale2 BScale2 - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); - a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); a_scale_thread_copy.Run(a_scale_grid_desc, a_scale_grid_buf, @@ -518,6 +501,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< c_thread_buf_per_scale; // Local prefetch A1 + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, 2, 1>{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { @@ -532,32 +516,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< }); }); -#if 0 - static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { - c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) - .template AsType()(Number{}) = 0; - }); - // Fill first mfma buffer - static_for<0, KRepeat, 1>{}([&](auto k0) { - vector_type a_thread_vec; - vector_type b_thread_vec; - - static_for<0, KPack, 1>{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = a_thread_buf - [Number{}]; - b_thread_vec.template AsType()(ik) = b_thread_bufs - [I0][Number{}]; - }); - - using mfma_input_type = - typename vector_type::type; - - xdlops_gemm.template Run<>(a_thread_vec.template AsType(), - b_thread_vec.template AsType(), - c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); - }); -#endif __builtin_amdgcn_sched_barrier(0); // main body @@ -574,8 +533,9 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< b_thread_bufs(local_read_buf)); b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(local_read_buf)); - a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + + block_sync_lds(); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(local_read_buf)); a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); a_scale_thread_copy.Run(a_scale_grid_desc, @@ -663,6 +623,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< if constexpr(m0.value == (MRepeat - 2)) { + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, KRepeat, 1>{}([&](auto k0) { @@ -766,7 +727,8 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< b_block_desc_n0_n1_k0_k1, b_block_origin_idx, b_thread_bufs(I1)); - a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(I1)); + block_sync_lds(); + a_blockwise_copy.Run(a_grid_desc, a_grid_buf, a_block_desc, a_block_buf.At(I1)); static_for<0, MRepeat, 1>{}([&](auto m0) { vector_type c_scale_thread_vec; @@ -820,6 +782,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3< if constexpr(m0.value == (MRepeat - 2)) { + __builtin_amdgcn_s_waitcnt(3952); block_sync_lds(); static_for<0, KRepeat, 1>{}([&](auto k0) { diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm_blockscale.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm_blockscale.hpp index 2257d3f690..01ce08e739 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm_blockscale.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm_blockscale.hpp @@ -1964,6 +1964,7 @@ struct GridwiseMoeGemmBlockScale BElementwiseOperation b_element_op, CElementwiseOperation c_element_op) { + ignore = a_element_op; ignore = b_element_op; const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1( IsInputGemm ? problem.NumTokens : problem.NumTokens * problem.TopK, @@ -2031,7 +2032,7 @@ struct GridwiseMoeGemmBlockScale constexpr auto AK1Threads = ABlockTransferThreadClusterLengths_AK0_M_AK1{}.At(I2); constexpr auto AKThreads = AK0Threads * AK1Threads; constexpr auto AMRepeats = MPerBlock / AMThreads; - const index_t token_pos = block_m_id * MPerBlock + threadIdx.x / AKThreads * AMRepeats; + const index_t token_pos = block_m_id * MPerBlock + threadIdx.x / AKThreads; if(token_pos >= max_token_id || expert_block_id * MPerBlock >= max_token_id || token0 >= problem.NumTokens) @@ -2039,7 +2040,7 @@ struct GridwiseMoeGemmBlockScale StaticallyIndexedArray gather_offsets; //= p_sorted_token_ids[token_pos]; static_for<0, AMRepeats, 1>{}([&](auto m0) { - const index_t fused_token = p_sorted_token_ids[token_pos + m0]; + const index_t fused_token = p_sorted_token_ids[token_pos + m0 * AMThreads]; index_t token_offset = fused_token & 0xffffff; if constexpr(!IsInputGemm) { @@ -2075,11 +2076,8 @@ struct GridwiseMoeGemmBlockScale // dummy constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(); // A matrix blockwise copy - auto a_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1_gather< + auto a_blockwise_copy = ThreadGroupTensorSliceTransfer_Gather_DirectLoad< ThisThreadBlock, - AElementwiseOperation, - ck::tensor_operation::element_wise::PassThrough, - InMemoryDataOperationEnum::Set, Sequence, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, @@ -2088,24 +2086,15 @@ struct GridwiseMoeGemmBlockScale decltype(a_grid_desc_ak0_m_ak1), decltype(a_block_desc_ak0_m_ak1), ABlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, ABlockTransferSrcVectorDim, 2, ABlockTransferSrcScalarPerVector, - ABlockTransferDstScalarPerVector_AK1, - 1, - 1, - AThreadTransferSrcResetCoordinateAfterRun, - true, IndexType, - 1, - BlockwiseGemmPipe::GlobalBufferNum>(a_grid_desc_ak0_m_ak1, - make_multi_index(0, 0, 0), - a_element_op, - a_block_desc_ak0_m_ak1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}, - gather_offsets); + 1>(a_grid_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + gather_offsets); // Thread-wise copy // K0 -> N0/NWave -> NWave -> KLane -> NLane -> KPack