From 1e87360ed1ca3628a8e54d7a958d147bcdcd4a80 Mon Sep 17 00:00:00 2001 From: valarLip <340077269@qq.com> Date: Thu, 15 May 2025 09:29:31 +0000 Subject: [PATCH] create reg spill debugging workspace --- .../65_gemm_multiply_multiply/CMakeLists.txt | 2 + .../moe_gemm1_xdl_fp8.cpp | 4 +- ..._pipeline_xdlops_b_preshuffle_selector.hpp | 23 +++++ .../gpu/grid/gridwise_moe_gemm.hpp | 85 +++++++++++++++++-- .../cpu/reference_moe_gemm.hpp | 8 +- 5 files changed, 110 insertions(+), 12 deletions(-) diff --git a/example/65_gemm_multiply_multiply/CMakeLists.txt b/example/65_gemm_multiply_multiply/CMakeLists.txt index ecf6079abb..0b430adf06 100644 --- a/example/65_gemm_multiply_multiply/CMakeLists.txt +++ b/example/65_gemm_multiply_multiply/CMakeLists.txt @@ -17,6 +17,8 @@ target_compile_options(example_gemm_multiply_multiply_xdl_fp8_bpreshuffle PRIVAT add_example_executable(example_moe_gemm1_xdl_fp8 moe_gemm1_xdl_fp8.cpp) add_example_executable(example_moe_gemm2_xdl_fp8 moe_gemm2_xdl_fp8.cpp) +add_definitions(-DDEBUG_SPILL) + list(APPEND gpu_list gfx942 gfx950) set(target 0) foreach(gpu IN LISTS GPU_TARGETS) 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 159fb21851..ce2d630450 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp @@ -157,8 +157,8 @@ using AElementOp = PassThrough; using BElementOp = PassThrough; static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; -static constexpr ck::index_t MPerBlock = 128; -static constexpr ck::index_t NPerBlock = 128; +static constexpr ck::index_t MPerBlock = 256; +static constexpr ck::index_t NPerBlock = 256; static constexpr ck::index_t MNPerXDL = 16; static constexpr ck::index_t MXDLPerWave = MPerBlock / (MNPerXDL * 1); static constexpr ck::index_t NXDLPerWave = NPerBlock / (MNPerXDL * 4); diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_selector.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_selector.hpp index df59513a88..404335bce4 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_selector.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_selector.hpp @@ -174,6 +174,7 @@ constexpr auto BlockGemmBPreshufflePipeline_Selector() { if constexpr(GUFusion) { +#ifdef DEBUG_SPILL return BlockwiseGemmXdlops_pipeline_bpreshuffle_gufusion_v3< BlkGemmPipeSche, BlockSize, @@ -195,6 +196,28 @@ constexpr auto BlockGemmBPreshufflePipeline_Selector() MRepeat, NRepeat, KPack>{}; +#else + return BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}; +#endif } else { 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 4c8424108d..3658d3c6df 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp @@ -2027,11 +2027,14 @@ struct GridwiseMoeGemm constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0); constexpr auto b_block_slice_copy_step = make_multi_index(0, 0, KRepeat, 0); + const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane( + (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) / + KPerBlock); + // Blockwise GEMM pipeline static_assert(std::is_default_constructible_v); auto blockwise_gemm_pipeline = BlockwiseGemmPipe{}; auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer(); - decltype(c_thread_buf) c_thread_buf_up; StaticBufferTupleOfVector c_thread_buf_fp32; + decltype(c_thread_buf) c_thread_buf_up; - const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane( - (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) / - KPerBlock); if constexpr(IsInputGemm) { +#ifdef DEBUG_SPILL const BDataType* p_b_grid_up = p_b_grid + expert_stride / 2 / BPackedSize; const auto b_grid_buf_up = make_dynamic_buffer( p_b_grid_up + expert_id * expert_stride / BPackedSize, @@ -2082,6 +2084,23 @@ struct GridwiseMoeGemm c_thread_buf, c_thread_buf_up, num_k_block_main_loop); +#else + (void)c_thread_buf_up; + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_bufs, + a_block_slice_copy_step, + b_grid_desc_bpreshuffled, + b_blockwise_copy, + b_grid_buf, + b_block_bufs, + b_block_slice_copy_step, + c_thread_buf, + num_k_block_main_loop); +#endif } else { @@ -2126,17 +2145,17 @@ struct GridwiseMoeGemm constexpr auto M2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I4); constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5); constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6); - constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); - // mul scales - const float* p_sorted_weights_0 = p_ds_grid[I0]; - const float* p_scale_b = p_ds_grid[I1]; static_assert(M0 * M1 * M2 * M3 * M4 == MPerBlock); static_assert(M4 == 4); const index_t m1 = get_warp_local_1d_id() / NWave; const index_t m3 = threadIdx.x % get_warp_size() / MPerXdl; - + constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); +#if 0 + // mul scales + const float* p_sorted_weights_0 = p_ds_grid[I0]; + const float* p_scale_b = p_ds_grid[I1]; if(p_sorted_weights_0 != nullptr && p_scale_b != nullptr) { if constexpr(PerTokenQuant) @@ -2306,7 +2325,55 @@ struct GridwiseMoeGemm }); }); } +#endif +#ifdef DEBUG_SPILL + vector_type topk_weights; // for gemm2 only + static_for<0, NXdlPerWave, 1>{}([&](auto n0) { + static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave + static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk + const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; + if constexpr(MulRoutedWeight) + { + topk_weights = *c_style_pointer_cast*>( + p_ds_grid[I2] + m_pos); + } + static_for<0, M4, 1>{}([&](auto m4) { // m_inst_group_size + constexpr index_t c_offset = + blockwise_gemm_pipeline.GetCThreadDesc().CalculateOffset( + make_tuple(m0, n0, m2 * M4 + m4)); + constexpr auto cidx = Number{}; + c_thread_buf_fp32(cidx) = c_thread_buf[cidx] + + c_thread_buf_up[cidx]; + }); + }); + }); + }); +#else + vector_type topk_weights; // for gemm2 only + static_for<0, NXdlPerWave, 1>{}([&](auto n0) { + static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave + static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk + const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; + if constexpr(MulRoutedWeight) + { + topk_weights = *c_style_pointer_cast*>( + p_ds_grid[I2] + m_pos); + } + static_for<0, M4, 1>{}([&](auto m4) { // m_inst_group_size + constexpr index_t c_offset = + blockwise_gemm_pipeline.GetCThreadDesc().CalculateOffset( + make_tuple(m0, n0, m2 * M4 + m4)); + constexpr auto cidx = Number{}; + + c_thread_buf_fp32(cidx) = c_thread_buf[cidx]; + }); + }); + }); + }); +#endif constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(); diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp index 120bf7484a..54c3c14385 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp @@ -169,7 +169,7 @@ struct ReferenceMoeGemm : public device::BaseOperator arg.c_element_op_(v_c, v_acc); arg.c_element_op_(v_c_up, v_acc_up); - +#if 0 if constexpr(ActivationType == 1) { v_c = v_c * arg.b_scale_e_n_(e, n) * arg.a_scale_t_(t); @@ -194,6 +194,12 @@ struct ReferenceMoeGemm : public device::BaseOperator v_c_up = v_c_up * arg.b_scale_e_n_(e, n + full_n) * arg.a_scale_t_(t); arg.c_t_k_n_(t, topk_id, n) = v_c * v_c_up; } +#endif +#ifdef DEBUG_SPILL + arg.c_t_k_n_(t, topk_id, n) = v_c + v_c_up; +#else + arg.c_t_k_n_(t, topk_id, n) = v_c; +#endif } };