From af2e0e51fc9ced7c595a5b93e783a12ce21d993b Mon Sep 17 00:00:00 2001 From: zanzhang Date: Sun, 27 Apr 2025 15:55:27 +0800 Subject: [PATCH] moe-gemm change into flatmm & tile_window modify to tile_scatter_gather --- example/ck_tile/XX_moe_gemm/moe_gemm.hpp | 6 +- .../ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp | 138 ++-- .../XX_moe_gemm/run_moe_gemm_example.inc | 128 ++- include/ck_tile/core.hpp | 1 + include/ck_tile/core/tensor/buffer_view.hpp | 2 +- include/ck_tile/core/tensor/load_tile.hpp | 87 +-- include/ck_tile/core/tensor/tensor_view.hpp | 32 +- .../core/tensor/tile_scatter_gather.hpp | 735 ++++++++++++++++++ include/ck_tile/core/tensor/tile_window.hpp | 28 +- .../core/tensor/tile_window_linear.hpp | 15 - .../ck_tile/core/tensor/tile_window_utils.hpp | 8 + .../ck_tile/core/tensor/transpose_tile.hpp | 108 +-- include/ck_tile/host/check_err.hpp | 12 +- .../reference_fused_single_moe_gemm.hpp | 1 + .../ops/epilogue/cshuffle_epilogue.hpp | 41 +- include/ck_tile/ops/gemm.hpp | 2 - .../ops/gemm/kernel/moe_gemm_kernel.hpp | 506 ------------ include/ck_tile/ops/moe_gemm.hpp | 15 + .../ops/moe_gemm/kernel/moe_gemm_kernel.hpp | 649 ++++++++++++++++ .../pipeline/moe_gemm_pipeline_ag_bg_cr.hpp | 0 .../moe_gemm_pipeline_ag_bg_cr_policy.hpp | 5 +- ..._gemm_pipeline_agmem_bgmem_creg_flatmm.hpp | 245 ++++++ ...ipeline_agmem_bgmem_creg_flatmm_policy.hpp | 37 + 23 files changed, 1957 insertions(+), 844 deletions(-) create mode 100644 include/ck_tile/core/tensor/tile_scatter_gather.hpp delete mode 100644 include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp create mode 100644 include/ck_tile/ops/moe_gemm.hpp create mode 100644 include/ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp rename include/ck_tile/ops/{gemm => moe_gemm}/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp (100%) rename include/ck_tile/ops/{gemm => moe_gemm}/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp (88%) create mode 100644 include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp create mode 100644 include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp diff --git a/example/ck_tile/XX_moe_gemm/moe_gemm.hpp b/example/ck_tile/XX_moe_gemm/moe_gemm.hpp index e7c647cb5c..1f2355f756 100644 --- a/example/ck_tile/XX_moe_gemm/moe_gemm.hpp +++ b/example/ck_tile/XX_moe_gemm/moe_gemm.hpp @@ -7,7 +7,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp" +#include "ck_tile/ops/moe_gemm.hpp" template struct GemmTypeConfig; @@ -37,6 +37,9 @@ auto create_args(int argc, char* argv[]) arg_parser.insert("experts", "8", "Num of experts - 8 by default") .insert("NumTokens", "128", "M dimensions - 128 by default.") .insert("TopK", "3", "Top K - 2 by default.") + // .insert("TopK", "2", "Top K - 2 by default.") + // .insert("N", "8192", "N dimensions - 4096 by default.") + // .insert("K", "6144", "K dimensions - 4096 by default.") .insert("N", "4096", "N dimensions - 4096 by default.") .insert("K", "4096", "K dimensions - 4096 by default.") .insert("stride_A", "", "Tensor A strides - it is empty by default.") @@ -46,6 +49,7 @@ auto create_args(int argc, char* argv[]) .insert("b_layout", "C", "B tensor data layout - Col by default.") .insert("c_layout", "R", "C tensor data layout - Row by default.") .insert("validate", "1", "0. No validation, 1. Validation on CPU.") + .insert("prec", "fp16", "data type. fp16/bf16/fp8/bf8") .insert("repeat", "10", "number of iterations to benchmark the kernel."); bool result = arg_parser.parse(argc, argv); diff --git a/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp b/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp index df000d633b..e34aa6159f 100644 --- a/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp +++ b/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp @@ -13,12 +13,11 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/flatmm.hpp" #include "ck_tile/host.hpp" #include "moe_gemm.hpp" #include "ck_tile/host/reference/reference_fused_single_moe_gemm.hpp" -namespace { - struct MoeGemmKernelParam { static const bool kPadM = false; @@ -30,8 +29,8 @@ struct MoeGemmKernelParam static const ck_tile::index_t N_Tile = 128; static const ck_tile::index_t K_Tile = 32; // need to ensure the M_per_thread = 1 - static const ck_tile::index_t M_Warp = 2; - static const ck_tile::index_t N_Warp = 2; + static const ck_tile::index_t M_Warp = 1; + static const ck_tile::index_t N_Warp = 4; static const ck_tile::index_t K_Warp = 1; static const ck_tile::index_t M_Warp_Tile = 32; @@ -39,85 +38,64 @@ struct MoeGemmKernelParam static const ck_tile::index_t K_Warp_Tile = 16; }; -using CodegenGemmShape = ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence>; - -using TilePartitioner = ck_tile::GemmTile1DPartitioner; - -template -using CodegenGemmTraits = ck_tile::TileGemmUniversalTraits; - -template -using CodegenPipelineProblem = - ck_tile::UniversalGemmPipelineProblem>; - -template -using CodegenGemmPipeline = - ck_tile::MoeGemmPipelineAgBgCrImpl>; - -template -using GemmEpilogue = ck_tile::CShuffleEpilogue::kBlockSize, - TilePartitioner::MPerBlock, - TilePartitioner::NPerBlock, - MoeGemmKernelParam::M_Warp, - MoeGemmKernelParam::N_Warp, - MoeGemmKernelParam::M_Warp_Tile, - MoeGemmKernelParam::N_Warp_Tile, - MoeGemmKernelParam::K_Warp_Tile, - CodegenPipelineProblem::TransposeC>>; - -// template -// using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue::TransposeC -// >>; - -template -using Kernel = ck_tile::MoeGemmKernel, - GemmEpilogue>; -}; // namespace - template float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s) { - using MoeGemmKernel = ::Kernel; + using CodegenMoeGemmShape = ck_tile::TileFlatmmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using CodegenMoeGemmTraits = ck_tile::TileGemmTraits; + + using CodegenPipelineProblem = + ck_tile::GemmPipelineProblem; + + using CodegenMoeGemmPolicy = ck_tile::UniversalFlatmmPipelineAgBgCrPolicy; + using CodegenMoeGemmPipeline = + ck_tile::MoeGemmPipelineAgBgCrImpl; + + using GemmEpilogue = ck_tile::CShuffleEpilogue>; + + using Kernel = ck_tile::MoeGemmKernel; // TODO: malloc sorted_tokend_ids buffer - const auto arguments = MoeGemmKernel::MoeGemmKernelArgs::MakeKernelArgs(gemm_desc); - - const dim3 grids = MoeGemmKernel::GridSize(gemm_desc.M, gemm_desc.N, 1); - constexpr dim3 blocks = MoeGemmKernel::BlockSize(); + const auto arguments = Kernel::MakeKernelArgs(gemm_desc); + const dim3 grids = Kernel::GridSize(gemm_desc.M, gemm_desc.N, 1); + constexpr dim3 blocks = Kernel::BlockSize(); // ck_tile::hip_check_error(hipMemcpyWithStream( // arguments.data(), @@ -127,7 +105,7 @@ float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s) if(s.log_level_ > 0) { - std::cout << "Launching kernel: " << MoeGemmKernel::GetName() << " with args:" + std::cout << "Launching kernel: " << Kernel::GetName() << " with args:" << " grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl; @@ -136,7 +114,7 @@ float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s) float ave_time = ck_tile::launch_kernel(s, ck_tile::make_kernel( - MoeGemmKernel{}, grids, blocks, 0, arguments)); + Kernel{}, grids, blocks, 0, arguments)); return ave_time; } diff --git a/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc b/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc index db049c9cc8..471742ab75 100644 --- a/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc +++ b/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc @@ -15,6 +15,40 @@ static constexpr inline auto is_row_major(Layout layout_) ck_tile::tensor_layout::gemm::RowMajor>>{}; } +template +auto shuffle_b(const ck_tile::HostTensor& t, std::string mfma_dtype, int mfma_type = 0) +{ + assert(t.get_lengths().size() == 2); + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + + if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 0) + { + ck_tile::HostTensor t_view({n_ / 32, 32, k_ / 16, 2, 8}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + } + else if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 1) + { + ck_tile::HostTensor t_view({n_ / 16, 16, k_ / 32, 4, 8}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + } + else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 0) + { + ck_tile::HostTensor t_view({n_ / 32, 32, k_ / 32, 2, 16}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + } + else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 1) + { + ck_tile::HostTensor t_view({n_ / 16, 16, k_ / 64, 4, 16}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + } + return t; +} + template auto calculate_rtol_atol(const ck_tile::index_t K, const ck_tile::index_t kbatch, @@ -37,7 +71,7 @@ auto calculate_rtol_atol(const ck_tile::index_t K, } template -float invoke_gemm(int n_warmup, int n_repeat, const moe_gemm_kargs& args) +float invoke_moe_gemm(int n_warmup, int n_repeat, const moe_gemm_kargs& args) { float ave_time = moe_gemm( args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); @@ -98,51 +132,69 @@ int run_moe_gemm_example_with_layouts(int argc, // TODO: replace the magic declaration const ck_tile::index_t MPerBlock = 128; + // const ck_tile::index_t max_num_tokens_padded = topk * num_tokens + experts * MPerBlock - topk; + // const ck_tile::index_t max_num_m_blocks = (max_num_tokens_padded + MPerBlock - 1) / MPerBlock; + ck_tile::index_t sorted_tile_num = 8; ck_tile::index_t valid_tile_num = sorted_tile_num; + // ck_tile::index_t sorted_tile_num = 16; + // ck_tile::index_t valid_tile_num = 13; + ck_tile::index_t sorted_size = sorted_tile_num * MPerBlock; + // ck_tile::index_t valid_size = valid_tile_num * MPerBlock; const ck_tile::index_t M = sorted_tile_num * MPerBlock; + // const ck_tile::index_t M = max_num_tokens_padded; std::unique_ptr a_m_k_dev_buf; - std::unique_ptr b_k_n_dev_buf; + std::unique_ptr b_origin_dev_buf; + std::unique_ptr b_shuffle_dev_buf; std::unique_ptr c_m_n_dev_buf; - stride_A = ck_tile::get_default_stride(M, N, stride_A, is_row_major(a_layout)); + stride_A = ck_tile::get_default_stride(num_tokens, K, stride_A, is_row_major(a_layout)); stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout)); - stride_C = ck_tile::get_default_stride(M, N, stride_C, is_row_major(CLayout{})); + stride_C = ck_tile::get_default_stride(num_tokens * topk, N, stride_C, is_row_major(CLayout{})); auto a_m_k_tensor = ck_tile::HostTensor( - ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout))); + ck_tile::host_tensor_descriptor(num_tokens, K, stride_A, is_row_major(a_layout))); // TODO: add the experts' weights in b auto b_k_n_tensor = ck_tile::HostTensor( is_row_major(b_layout) - ? ck_tile::host_tensor_descriptor(experts * K, N, stride_B, is_row_major(b_layout)) + ? ck_tile::host_tensor_descriptor(experts * N, K, stride_B, is_row_major(b_layout)) : ck_tile::host_tensor_descriptor(K, experts * N, stride_B, is_row_major(b_layout))); - auto c_m_n_tensor = ck_tile::HostTensor( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); - std::cout << "gemm" - << " a_m_k: " << a_m_k_tensor.mDesc << " b_k_n: " << b_k_n_tensor.mDesc - << " c_m_n: " << c_m_n_tensor.mDesc << std::endl; + std::string mfma = arg_parser.get_str("prec"); + + auto c_m_n_tensor = ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(num_tokens * topk, N, stride_C, is_row_major(CLayout{}))); ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k_tensor); ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n_tensor); + auto b_shuffle_host = shuffle_b(b_k_n_tensor, mfma, 0); + + std::cout << "gemm" + << " a_m_k: " << a_m_k_tensor.mDesc << " b_k_n: " << b_k_n_tensor.mDesc + << " b_shuffle: " << b_shuffle_host.mDesc << " c_m_n: " << c_m_n_tensor.mDesc << std::endl; + a_m_k_dev_buf = std::make_unique(a_m_k_tensor.get_element_space_size_in_bytes()); - b_k_n_dev_buf = + b_origin_dev_buf = std::make_unique(b_k_n_tensor.get_element_space_size_in_bytes()); + b_shuffle_dev_buf = + std::make_unique(b_shuffle_host.get_element_space_size_in_bytes()); c_m_n_dev_buf = std::make_unique(c_m_n_tensor.get_element_space_size_in_bytes()); a_m_k_dev_buf->ToDevice(a_m_k_tensor.data()); - b_k_n_dev_buf->ToDevice(b_k_n_tensor.data()); + b_origin_dev_buf->ToDevice(b_k_n_tensor.data()); + b_shuffle_dev_buf->ToDevice(b_shuffle_host.data()); c_m_n_dev_buf->SetZero(); c_m_n_tensor.SetZero(); const void* p_a = a_m_k_dev_buf->GetDeviceBuffer(); - const void* p_b = b_k_n_dev_buf->GetDeviceBuffer(); + const void* p_b_origin = b_origin_dev_buf->GetDeviceBuffer(); + const void* p_b_shuffle = b_shuffle_dev_buf->GetDeviceBuffer(); void* p_c = c_m_n_dev_buf->GetDeviceBuffer(); // TODO: malloc and init sorted tokens and max tokens buffer @@ -150,7 +202,7 @@ int run_moe_gemm_example_with_layouts(int argc, ck_tile::HostTensor expert_ids( ck_tile::HostTensorDescriptor({sorted_tile_num}, {1})); ck_tile::HostTensor sorted_token_ids( - ck_tile::HostTensorDescriptor({sorted_tile_num * MPerBlock}, {1})); + ck_tile::HostTensorDescriptor({sorted_size}, {1})); ck_tile::HostTensor max_token_id( ck_tile::HostTensorDescriptor({1 + sorted_tile_num})); @@ -163,11 +215,14 @@ int run_moe_gemm_example_with_layouts(int argc, max_token_id.mData = {valid_tile_num * MPerBlock, 0, 1, 2, 3, 4, 6, 7, 8, 8}; int eids[] = {0, 1, 2, 3, 4, 4, 5, 6, 3, 3, 3, 3}; // {2, 1, 1, 2, 2, 2, 1, 2} + // max_token_id.mData = {valid_size, 2, 2, 1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 0, 0, 0}; + // int eids[] = {0, 0, 1, 2, 3, 3, 4, 4, 5, 5, 6, 7, 7, 3, 3, 3}; for(int i = 0; i < sorted_tile_num; i++) { expert_ids.mData[i] = eids[i]; } - int token_per_tile = (num_tokens * topk + valid_tile_num - 1) / valid_tile_num; + // int token_per_tile = (num_tokens * topk + valid_tile_num - 1) / valid_tile_num; + int token_per_tile = num_tokens * topk / valid_tile_num; int tokenid = 0; // sorted_token_ids.mData[0] = 0; for(int i = 0; i < sorted_tile_num * MPerBlock; i++) @@ -199,10 +254,11 @@ int run_moe_gemm_example_with_layouts(int argc, p_expert_ids_dev, p_max_token_id_dev, p_a, - p_b, + p_b_shuffle, p_c, num_tokens, topk, + 1, //k_batch M, N, K, @@ -210,7 +266,7 @@ int run_moe_gemm_example_with_layouts(int argc, stride_B, stride_C}; - invoke_gemm(3, repeat, gemm_desc); + invoke_moe_gemm(3, repeat, gemm_desc); c_m_n_dev_buf->FromDevice(c_m_n_tensor.data()); @@ -218,7 +274,7 @@ int run_moe_gemm_example_with_layouts(int argc, if(arg_parser.get_int("validate")) { ck_tile::HostTensor c_m_n_host_ref( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); + ck_tile::host_tensor_descriptor(num_tokens * topk, N, stride_C, is_row_major(CLayout{}))); c_m_n_host_ref.SetZero(); @@ -238,7 +294,7 @@ int run_moe_gemm_example_with_layouts(int argc, p_expert_ids_dev, p_max_token_id_dev, static_cast(p_a), - static_cast(p_b), + static_cast(p_b_origin), static_cast(c_m_n_ref_buf->GetDeviceBuffer()), num_tokens, topk, @@ -254,22 +310,22 @@ int run_moe_gemm_example_with_layouts(int argc, K, 1 /*kbatch*/, max_accumulated_value); c_m_n_ref_buf->FromDevice(c_m_n_host_ref.data()); - for(int im = 0; im < M; im++) - { - for(int in = 0; in < N; in++) - { - // if (static_cast(static_cast(p_c)[im * N + in]) != 0) - printf("c[%d][%d]: %f ", - im, - in, - static_cast(static_cast(p_c)[im * N + in])); - printf("ref[%d][%d]: %f \n", - im, - in, - static_cast( - static_cast(c_m_n_host_ref.data())[im * N + in])); - } - } + // for(int im = 0; im < M; im++) + // { + // for(int in = 0; in < N; in++) + // { + // // if (static_cast(static_cast(p_c)[im * N + in]) != 0) + // printf("c[%d][%d]: %f ", + // im, + // in, + // static_cast(static_cast(p_c)[im * N + in])); + // printf("ref[%d][%d]: %f \n", + // im, + // in, + // static_cast( + // static_cast(c_m_n_host_ref.data())[im * N + in])); + // } + // } pass = ck_tile::check_err(c_m_n_tensor, c_m_n_host_ref, diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index 821b3a8e84..c341a20c0b 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -55,6 +55,7 @@ #include "ck_tile/core/tensor/tile_elementwise.hpp" #include "ck_tile/core/tensor/tile_window.hpp" #include "ck_tile/core/tensor/tile_window_linear.hpp" +#include "ck_tile/core/tensor/tile_scatter_gather.hpp" #include "ck_tile/core/tensor/tile_window_utils.hpp" #include "ck_tile/core/tensor/transpose_tile.hpp" #include "ck_tile/core/tensor/update_tile.hpp" diff --git a/include/ck_tile/core/tensor/buffer_view.hpp b/include/ck_tile/core/tensor/buffer_view.hpp index bdcfbdd920..c7e24cbc2b 100644 --- a/include/ck_tile/core/tensor/buffer_view.hpp +++ b/include/ck_tile/core/tensor/buffer_view.hpp @@ -5,7 +5,7 @@ #include "ck_tile/core/config.hpp" #include "ck_tile/core/arch/arch.hpp" -#if __clang_major__ == 20 +#if __clang_major__ >= 20 #include "ck_tile/core/arch/amd_buffer_addressing_builtins.hpp" #else #include "ck_tile/core/arch/amd_buffer_addressing.hpp" diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index b280a1725d..0212e773f0 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -18,32 +18,10 @@ namespace ck_tile { -template -CK_TILE_DEVICE auto load_tile(const tile_window_with_static_distribution& tile_window, - number = {}, - bool_constant = {}) -{ - return tile_window.load(number{}, bool_constant{}); -} - -template -CK_TILE_DEVICE auto load_tile(const tile_window_linear& tile_window, +CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window, number = {}, bool_constant = {}) { @@ -51,35 +29,11 @@ CK_TILE_DEVICE auto load_tile(const tile_window_linear CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile, - const tile_window_with_static_distribution& tile_window, - number = {}, - bool_constant = {}) -{ - return tile_window.load(dst_tile, number{}, bool_constant{}); -} - -template -CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile, - const tile_window_linear& tile_window, + const TileWindow_& tile_window, number = {}, bool_constant = {}) { @@ -138,19 +92,13 @@ CK_TILE_DEVICE auto load_tile_raw(T& tile, } template CK_TILE_DEVICE auto async_load_tile_raw(LdsTileWindow_&& lds_tile, - const tile_window_with_static_distribution& tile_window, + const TileWindow_& tile_window, number = {}, bool_constant = {}, bool_constant = {}) @@ -161,29 +109,6 @@ async_load_tile_raw(LdsTileWindow_&& lds_tile, bool_constant{}); } -template -CK_TILE_DEVICE auto async_load_tile_raw(LdsTileWindow_&& lds_tile, - const tile_window_linear& tile_window, - number = {}, - bool_constant = {}, - bool_constant = {}) -{ - return tile_window.async_load_raw(lds_tile, - number{}, - bool_constant{}, - bool_constant{}); -} - CK_TILE_DEVICE auto async_load_fence(index_t cnt = 0) { asm volatile("s_waitcnt vmcnt(%0)" : : "n"(cnt) : "memory"); diff --git a/include/ck_tile/core/tensor/tensor_view.hpp b/include/ck_tile/core/tensor/tensor_view.hpp index 32de227b52..b1b1f36bb9 100644 --- a/include/ck_tile/core/tensor/tensor_view.hpp +++ b/include/ck_tile/core/tensor/tensor_view.hpp @@ -209,6 +209,26 @@ struct tensor_view coordinate_has_valid_offset_assuming_top_index_is_valid(desc_, coord), bool_constant{}); } + template >::scalar_type, + typename vector_traits>::scalar_type>, + bool>::type = false> + CK_TILE_HOST_DEVICE constexpr void + async_get_vectorized_elements_raw(remove_cvref_t* smem, + const TensorCoord& coord, + index_t coord_extra_offset, + index_t linear_offset, + bool_constant = {}) const + { + return buf_.template async_get_raw( + smem, + (coord.get_offset() + coord_extra_offset) / PackedSize, + linear_offset / PackedSize, + coordinate_has_valid_offset_assuming_top_index_is_valid(desc_, coord), + bool_constant{}); + } template CK_TILE_HOST_DEVICE constexpr auto make_tensor_view(DataType* p, const tensor_descriptor& desc) { - auto buffer_view = - make_buffer_view(p, desc.get_element_space_size()); + auto buffer_view = make_buffer_view(p, desc.get_element_space_size()); return tensor_view{buffer_view, desc}; } template {}, number{}); - auto buffer_view = - make_buffer_view(p, desc.get_element_space_size()); + auto buffer_view = make_buffer_view(p, desc.get_element_space_size()); return tensor_view{buffer_view, desc}; } template @@ -463,8 +478,7 @@ make_naive_tensor_view_packed(DataType* p, auto desc = make_naive_tensor_descriptor_packed(lengths, number{}); - auto buffer_view = - make_buffer_view(p, desc.get_element_space_size()); + auto buffer_view = make_buffer_view(p, desc.get_element_space_size()); return tensor_view{buffer_view, desc}; } diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp new file mode 100644 index 0000000000..280ffd8c5f --- /dev/null +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -0,0 +1,735 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core/arch/utility.hpp" +#include "ck_tile/core/algorithm/space_filling_curve.hpp" +#include "ck_tile/core/config.hpp" +#include "ck_tile/core/container/array.hpp" +#include "ck_tile/core/container/sequence.hpp" +#include "ck_tile/core/container/tuple.hpp" +#include "ck_tile/core/container/container_helper.hpp" +#include "ck_tile/core/tensor/static_distributed_tensor.hpp" +#include "ck_tile/core/tensor/tensor_adaptor.hpp" +#include "ck_tile/core/tensor/tile_distribution.hpp" +#include "ck_tile/core/utility/functional.hpp" +#include "ck_tile/core/utility/type_traits.hpp" + +namespace ck_tile { + +/** + * @brief This class provides tile (windowed) view and access to the device memory. + * + * @note This tile window does not support single issue you need to use tile_window_linear + * structure for this purpose + * + * @tparam BottomTensorView_ Class describing & holding device tensor memory. + * @tparam WindowLengths_ Spatial sizes of windowed view on tensor. + * @tparam StaticTileDistribution_ Thread distribution (mapping) into Tile dimensions + * @tparam NumCoord TBD + */ +template +struct tile_scatter_gather +{ + using BottomTensorView = remove_reference_t; + using WindowLengths = remove_cvref_t; + using TileDstr = remove_cvref_t; + using PageIdxArray = remove_cvref_t; + using WindowAdaptor = typename TileDstr::PsYs2XsAdaptor; + using BottomTensorDesc = typename BottomTensorView::TensorDesc; + + using DataType = remove_cvref_t; + + static constexpr index_t NDimWindowAdaptorTop = WindowAdaptor::get_num_of_top_dimension(); + static constexpr index_t NDimBottomTensor = BottomTensorDesc::get_num_of_dimension(); + + static constexpr index_t NDimP = TileDstr::get_num_of_dimension_p(); + static constexpr index_t NDimY = TileDstr::get_num_of_dimension_y(); + + static constexpr auto I0 = number<0>{}; + static constexpr auto I1 = number<1>{}; + static_assert(NumCoord == 1); + + // TODO: check WindowLengths and StaticTileDistribution are consistent + + static_assert(ck_tile::is_known_at_compile_time::value, + "wrong! lengths should be static"); + static_assert(TileDstr::is_static(), "wrong!"); + + static_assert(NDimBottomTensor == WindowAdaptor::get_num_of_bottom_dimension(), + "wrong! inconsistent # of diemsnions"); + + using AdaptorTopIndex = array; + using BottomTensorIndex = array; + + using WindowAdaptorCoord = + decltype(make_tensor_adaptor_coordinate(WindowAdaptor{}, AdaptorTopIndex{})); + + using BottomTensorCoord = + decltype(make_tensor_coordinate(BottomTensorDesc{}, BottomTensorIndex{})); + + struct load_store_traits + { + private: + static constexpr auto get_vector_dim_y_scalar_per_vector() + { + const auto [ys_vector_lengths, ys_vector_strides] = + tile_scatter_gather:: + get_window_adaptor_ys_safe_vector_length_strides(); + + index_t VectorDimY_ = 0; + index_t ScalarPerVector_ = 1; + + for(index_t i = 0; i < NDimY; ++i) + { + if(ys_vector_strides[i] == 1 && ys_vector_lengths[i] > ScalarPerVector_) + { + ScalarPerVector_ = ys_vector_lengths[i]; + VectorDimY_ = i; + } + } + + return make_tuple(VectorDimY_, ScalarPerVector_); + } + + public: + static constexpr index_t PackedSize = + ck_tile::numeric_traits>::PackedSize; + static constexpr index_t VectorDimY = get_vector_dim_y_scalar_per_vector().template at<0>(); + static constexpr index_t ScalarPerVector = + get_vector_dim_y_scalar_per_vector().template at<1>(); + + // using vector_type_t = vector_type_maker_t; + // using vector_t = typename vector_type_t::type; + using vector_t = thread_buffer; + + private: + static constexpr auto scalars_per_access_ = [] { + constexpr auto scalars_per_access_arr = generate_array( + [&](auto i) { return (i == VectorDimY) ? ScalarPerVector : 1; }, number{}); + + /// TODO: add non-automatic storage argument support to macro TO_SEQUENCE() + constexpr auto NDimY_ = NDimY; + + return TO_SEQUENCE(scalars_per_access_arr, NDimY_); + }(); + + static constexpr auto get_space_filling_curve() + { + constexpr auto tile_dstr = TileDstr{}; + + constexpr auto thread_tensor_lengths_ys = + to_sequence(tile_dstr.get_ys_to_d_descriptor().get_lengths()); + + // FIXME: need logic to judge dim access order + using DimAccessOrder = typename arithmetic_sequence_gen<0, NDimY, 1>::type; + + return space_filling_curve{}; + } + + public: + using SFC_Ys = decltype(get_space_filling_curve()); + + static constexpr index_t NumAccess = SFC_Ys::get_num_of_access(); + + static_assert(0 < NumAccess, "Wrong! NumAccess should be larger than 0"); + static_assert(NumAccess % NumCoord == 0, "wrong! # of access is not divisible by NumCoord"); + }; + + static constexpr index_t NumAccessPerCoord = load_store_traits::NumAccess / NumCoord; + + CK_TILE_DEVICE constexpr tile_scatter_gather() = default; + + CK_TILE_DEVICE constexpr tile_scatter_gather( + const BottomTensorView& bottom_tensor_view, + const WindowLengths& window_lengths, + const BottomTensorIndex& window_origin, + const TileDstr& tile_distribution, + const PageIdxArray& page_idx) + : bottom_tensor_view_{bottom_tensor_view}, + window_lengths_{window_lengths}, + window_origin_{window_origin}, + tile_dstr_{tile_distribution}, + page_idx_{page_idx}, + pre_computed_coords_{} + { +#if 0 // debug + // TODO: this use more register for FA, but less register for GEMM + // need investigation + // only support warp-tile and block-tile + static_assert(NDimP == 1 or NDimP == 2, "wrong!"); + + WindowAdaptorCoord window_adaptor_thread_coord_tmp; + + if constexpr(NDimP == 1) + { + window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate( + tile_distribution.get_ps_ys_to_xs_adaptor(), AdaptorTopIndex{get_lane_id(), 0}); + } + else if constexpr(NDimP == 2) + { + window_adaptor_thread_coord_tmp = + make_tensor_adaptor_coordinate(tile_distribution.get_ps_ys_to_xs_adaptor(), + AdaptorTopIndex{get_warp_id(), get_lane_id(), 0}); + } +#else + // TODO: this use less register for FA, but more register for GEMM + // need investigation + const auto window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate( + tile_distribution.get_ps_ys_to_xs_adaptor(), + container_concat(detail::get_partition_index(tile_distribution), + array{0})); +#endif + + BottomTensorIndex bottom_tensor_thread_origin_idx_tmp = + window_origin + window_adaptor_thread_coord_tmp.get_bottom_index(); + bottom_tensor_thread_origin_idx_tmp(HsGatherDim) = 0; + const auto bottom_tensor_thread_coord_tmp = make_tensor_coordinate( + bottom_tensor_view_.get_tensor_descriptor(), bottom_tensor_thread_origin_idx_tmp); + + // pre-compute NumCoord (WindowAdaptorCoord, BottomTensorCoord) bundles to speed up + // future load/store() calls (might allocate more registers) + using Traits = load_store_traits; + using SFC_Ys = typename Traits::SFC_Ys; + + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + auto window_adaptor_thread_coord = window_adaptor_thread_coord_tmp; + auto bottom_tensor_thread_coord = bottom_tensor_thread_coord_tmp; + + constexpr auto idx_diff_ys = + SFC_Ys::get_step_between(number<0>{}, number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), idx_diff_ys); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + + pre_computed_coords_(iCoord) = + make_tuple(window_adaptor_thread_coord, bottom_tensor_thread_coord); + }); + } + + CK_TILE_DEVICE static constexpr index_t get_num_of_dimension() { return NDimBottomTensor; } + + CK_TILE_DEVICE static constexpr bool has_static_tile_distribution() + { + return TileDstr::is_static(); + } + + CK_TILE_DEVICE constexpr auto get_window_lengths() const { return window_lengths_; } + + CK_TILE_DEVICE constexpr auto get_tile_distribution() const { return tile_dstr_; } + + CK_TILE_DEVICE constexpr auto get_bottom_tensor_view() const { return bottom_tensor_view_; } + + CK_TILE_DEVICE constexpr auto get_window_origin() const { return window_origin_; } + + CK_TILE_DEVICE constexpr void + set_bottom_tensor_view_data_ptr(typename BottomTensorView::DataType* data) + { + bottom_tensor_view_.buf_.p_data_ = data; + } + + // move thread's window adaptor coordinate and bottom tensor coordinate + // [p0, p1, ..., y0, y1, ...] ==> [x0, x1, ...] ==> [x0', x1', ...] ==> [offset] + template + CK_TILE_DEVICE void move_window_adaptor_and_bottom_tensor_thread_coordinate( + WindowAdaptorCoord& window_adaptor_thread_coord, + BottomTensorCoord& bottom_tensor_thread_coord, + const ATopIndex& idx_diff_adaptor_top) const + { + array idx_diff_adaptor_bottom; + + move_tensor_adaptor_coordinate(tile_dstr_.get_ps_ys_to_xs_adaptor(), + window_adaptor_thread_coord, + idx_diff_adaptor_top, + idx_diff_adaptor_bottom); + + move_tensor_coordinate(bottom_tensor_view_.get_tensor_descriptor(), + bottom_tensor_thread_coord, + idx_diff_adaptor_bottom); + } + + // return vector dimension among [y0, y1, ...] + CK_TILE_DEVICE static constexpr auto get_window_adaptor_ys_safe_vector_length_strides() + { + // bottom tensor top dimension vector lengths and strides + const auto [bottom_tensor_top_dim_vector_lengths, bottom_tensor_top_dim_vector_strides] = + BottomTensorDesc::get_top_dimension_safe_vector_length_strides(); + + // window vector lengths/strides + const auto window_adaptor_bottom_dim_vector_lengths = bottom_tensor_top_dim_vector_lengths; + const auto window_adaptor_bottom_dim_vector_strides = bottom_tensor_top_dim_vector_strides; + + // window adaptor [p0, p1, ..., y0, y1, ...] + array window_adaptor_vector_lengths{ + -1}; + array window_adaptor_vector_strides{ + -1}; + + constexpr auto window_adaptor_bottom_dims = + WindowAdaptor::get_bottom_dimension_hidden_ids(); + + set_container_subset(window_adaptor_vector_lengths, + window_adaptor_bottom_dims, + window_adaptor_bottom_dim_vector_lengths); + set_container_subset(window_adaptor_vector_strides, + window_adaptor_bottom_dims, + window_adaptor_bottom_dim_vector_strides); + + const auto [window_adaptor_ps_ys_vector_lengths, window_adaptor_ps_ys_vector_strides] = + WindowAdaptor{}.get_top_dimension_safe_vector_length_strides( + window_adaptor_vector_lengths, window_adaptor_vector_strides); + + // [y0, y1, ...] + constexpr auto y_dims = typename arithmetic_sequence_gen::type{}; + + return make_tuple(get_container_subset(window_adaptor_ps_ys_vector_lengths, y_dims), + get_container_subset(window_adaptor_ps_ys_vector_strides, y_dims)); + } + + CK_TILE_DEVICE constexpr auto get_num_of_access() const { return load_store_traits::NumAccess; } + + template + CK_TILE_DEVICE auto load(number = {}, + bool_constant = {}) const + { + constexpr auto tile_dstr = TileDstr{}; + auto dst_tensor = make_static_distributed_tensor(tile_dstr); + load(dst_tensor, number{}, bool_constant{}); + return dst_tensor; + } + + template + CK_TILE_DEVICE auto load(DistributedTensor& dst_tensor, + number = {}, + bool_constant = {}) const + { + using Traits = load_store_traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number{}]; + const auto page_offset = page_idx_[idx_gather]; + // read from bottom tensor + const vector_t vec_value = + get_bottom_tensor_view().template get_vectorized_elements( + bottom_tensor_thread_coord, page_offset, bool_constant{}); +#if 1 + // write into distributed tensor + static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { + constexpr auto idx_ys = generate_tuple( + [&](auto jj) { + return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j) + : idx_ys_start[jj]; + }, + number{}); + + constexpr index_t d = + tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) / + Traits::PackedSize; + + dst_tensor.get_thread_buffer().template at() = + vec_value.template get_as()[j / Traits::PackedSize]; + }); +#else + constexpr index_t d = + tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start); + static_assert(d % Traits::ScalarPerVector == 0); + + dst_tensor.get_thread_buffer().template get_as()( + number{}) = bit_cast(vec_value); +#endif + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + } + }); + }); + } + + + // TODO: currently async load only implemented in inline asm + template + CK_TILE_DEVICE auto async_load_raw(LdsTileWindow_&& lds_tile, + number = {}, + bool_constant = {}, + bool_constant = {}) const + { + using LdsTileWindow = remove_cvref_t; + // using LdsTensorView = typename LdsTileWindow::BottomTensorView; + using LdsDataType = typename LdsTileWindow::DataType; + // using LdsDescriptor = typename LdsTileWindow::BottomTensorDesc; + + // issues * warps * lanes + static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded + + const index_t size_per_buf = + lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( + make_tuple(number<0>{}, number<0>{}, number<0>{})) * + sizeof(LdsDataType); + + const index_t size_per_wave = + lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( + make_tuple(number<0>{}, number<1>{}, number<0>{})) * + sizeof(LdsDataType) - + size_per_buf; + + const index_t size_per_issue = + lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( + make_tuple(number<1>{}, number<0>{}, number<0>{})) * + sizeof(LdsDataType) - + size_per_buf; + + const index_t m0_init_value = size_per_buf + size_per_wave * get_warp_id(); + m0_set_with_memory(m0_init_value); // This should be wave independent + + using Traits = load_store_traits; + + // using vector_type_t = typename Traits::vector_type_t; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + LdsDataType* smem = lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_; + + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + constexpr auto pre_nop_ = [&]() { + if constexpr(pre_nop && iCoord == 0 && iCoordAccess == 0) + return bool_constant{}; + else + return bool_constant{}; + }(); + + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number{}]; + const auto page_offset = page_idx_[idx_gather]; + // read from bottom tensor + get_bottom_tensor_view().template async_get_vectorized_elements_raw( + smem, bottom_tensor_thread_coord, page_offset, 0, pre_nop_); + + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + + m0_inc_with_memory(size_per_issue); + } + }); + }); + } + + template + CK_TILE_DEVICE void store(const static_distributed_tensor& dstr_tensor, + number = {}, + bool_constant = {}) const + { + using Traits = load_store_traits; + + // using vector_type_t = typename Traits::vector_type_t; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + // printf("off %d\n", page_idx_[I0]); + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number<0>{}]; + const auto page_offset = page_idx_[idx_gather]; + + // printf("idx_ys_start[0], idx_ys_start[1](%d, %d) \n", + // idx_ys_start[number<0>{}]+0, idx_ys_start[number<1>{}]+0); + + // read from distributed tensor + // vector_type_t vec; + vector_t vec_value; + + static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { + constexpr auto idx_ys = generate_tuple( + [&](auto jj) { + return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j) + : idx_ys_start[jj]; + }, + number{}); + + constexpr index_t d = + tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) / + Traits::PackedSize; + // printf("thread_idx_m: %d j: %d\n", idx_ys[number<0>{}] + 0, 0+j); + vec_value.template get_as()(j / Traits::PackedSize) = + dstr_tensor.get_thread_buffer().template at(); + }); + + // const vector_t vec_value = vec.template get_as().template at<0>(); + + // write into bottom tensor + get_bottom_tensor_view().template set_vectorized_elements( + bottom_tensor_thread_coord, + page_offset, + vec_value, + bool_constant{}); + // printf("coord_offset:%d, scatter_offset:%d \n", + // bottom_tensor_thread_coord.get_offset(), offset); move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + } + }); + }); + } + + // move thread's botom tensor coordiante + // [x0', x1', ... ] ==> [offset] + // also move window-origin + CK_TILE_DEVICE void move(const BottomTensorIndex& step) + { + window_origin_ += step; + BottomTensorIndex step_new = step; + step_new(HsGatherDim) = 0; + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + move_tensor_coordinate(bottom_tensor_view_.get_tensor_descriptor(), + pre_computed_coords_(iCoord)(I1), + step_new); + }); + } + + CK_TILE_DEVICE void update_page_idx(const PageIdxArray& new_idx) + { + page_idx_ = new_idx; + + // static_for<0, 2, 1>{}([&](auto k0) { + // printf("update tid %d %d \n", threadIdx.x, page_idx_[k0]); + // }); + } + CK_TILE_DEVICE void set_window_origin(const BottomTensorIndex& new_window_origin) + { + window_origin_ = new_window_origin; + +#if 0 // debug + // TODO: this use more register for FA, but less register for GEMM + // need investigation + // only support warp-tile and block-tile + static_assert(NDimP == 1 or NDimP == 2, "wrong!"); + + WindowAdaptorCoord window_adaptor_thread_coord_tmp; + + if constexpr(NDimP == 1) + { + window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate( + tile_dstr_.get_ps_ys_to_xs_adaptor(), AdaptorTopIndex{get_lane_id(), 0}); + } + else if constexpr(NDimP == 2) + { + window_adaptor_thread_coord_tmp = + make_tensor_adaptor_coordinate(tile_dstr_.get_ps_ys_to_xs_adaptor(), + AdaptorTopIndex{get_warp_id(), get_lane_id(), 0}); + } +#else + // TODO: this use less register for FA, but more register for GEMM + // need investigation + const auto window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate( + tile_dstr_.get_ps_ys_to_xs_adaptor(), + container_concat(detail::get_partition_index(tile_dstr_), array{0})); +#endif + + BottomTensorIndex bottom_tensor_thread_origin_idx_tmp = + window_origin_ + window_adaptor_thread_coord_tmp.get_bottom_index(); + + bottom_tensor_thread_origin_idx_tmp(HsGatherDim) = 0; + const auto bottom_tensor_thread_coord_tmp = make_tensor_coordinate( + bottom_tensor_view_.get_tensor_descriptor(), bottom_tensor_thread_origin_idx_tmp); + + // pre-compute NumCoord (WindowAdaptorCoord, BottomTensorCoord) bundles to speed up + // future load/store() calls (might allocate more registers) + using Traits = load_store_traits; + using SFC_Ys = typename Traits::SFC_Ys; + + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + auto window_adaptor_thread_coord = window_adaptor_thread_coord_tmp; + auto bottom_tensor_thread_coord = bottom_tensor_thread_coord_tmp; + + constexpr auto idx_diff_ys = + SFC_Ys::get_step_between(number<0>{}, number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), idx_diff_ys); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + + pre_computed_coords_(iCoord) = + make_tuple(window_adaptor_thread_coord, bottom_tensor_thread_coord); + }); + } + + CK_TILE_HOST_DEVICE void init_raw() { bottom_tensor_view_.init_raw(); } + + // this is the bottom tensor view + // [x0', x1', ...] ==> [offset] + BottomTensorView bottom_tensor_view_; + + // + WindowLengths window_lengths_; + + // origin ([x0', x1', ...]) of window on bottom tensor + BottomTensorIndex window_origin_; + + // Tile tensor distribution, which contains: + // 1. adaptor for window: [p0, p1, ..., y0, y1, ...] ==> [x0, x1, ...] + // 2. thread descriptor for thread tensor in register: [y0, y1, ...] ==> [d] + TileDstr tile_dstr_; + + PageIdxArray page_idx_; + + // this contains: + // per-thread coordinate for window adaptor + // per-thread coordinate for bottom tensor + array, NumCoord> pre_computed_coords_; +}; + +// TODO: use strategy +template +CK_TILE_DEVICE constexpr auto +make_tile_scatter_gather(const TensorView_& tensor_view, + const WindowLengths_& window_lengths, + const multi_index& origin, + const StaticTileDistribution_& tile_distribution, + const StaticPageIndexArray_& page_idx, + number = {}, + number = {}) +{ + return tile_scatter_gather, + remove_cvref_t, + remove_cvref_t, + remove_cvref_t, + HsGatherDim, + NumCoord>{ + tensor_view, window_lengths, origin, tile_distribution, page_idx}; +} + +template +CK_TILE_DEVICE constexpr auto +make_tile_scatter_gather(const tile_window_with_static_lengths& tile_window, + const multi_index& origin, + const StaticTileDistribution& tile_distribution, + const StaticPageIndexArray& page_idx, + number = {}) +{ + return make_tile_scatter_gather(tile_window.get_bottom_tensor_view(), + tile_window.get_window_lengths(), + origin, + tile_distribution, + page_idx, + number{}); +} + +template +CK_TILE_DEVICE constexpr auto +make_tile_scatter_gather(const tile_window_with_static_lengths& tile_window, + const StaticTileDistribution& tile_distribution, const StaticPageIndexArray& page_idx, + number = {}) +{ + return make_tile_scatter_gather(tile_window.get_bottom_tensor_view(), + tile_window.get_window_lengths(), + tile_window.get_window_origin(), + tile_distribution, + page_idx, + number{}); +} + +// template +// CK_TILE_DEVICE constexpr auto +// make_tile_window_raw(const tile_window_with_static_lengths& tile_window, +// const StaticTileDistribution& tile_distribution) +// { +// auto w = make_tile_scatter_gather(tile_window.get_bottom_tensor_view(), +// tile_window.get_window_lengths(), +// tile_window.get_window_origin(), +// tile_distribution); +// w.init_raw(); +// return w; +// } + + +} // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index eac098a184..9faa58f03c 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -613,7 +613,7 @@ struct tile_window_with_static_distribution index_t i_access_unsupport_ = -1, bool oob_conditional_check = true> CK_TILE_DEVICE void store(const static_distributed_tensor& dstr_tensor, - const statically_indexed_array offsets, + const statically_indexed_array& offsets, number = {}, bool_constant = {}) const { @@ -1097,23 +1097,6 @@ make_tile_window_raw(const TensorView_& tensor_view, return w; } -template -CK_TILE_DEVICE void move_tile_window( - tile_window_with_static_distribution& window, - const typename tile_window_with_static_distribution::BottomTensorIndex& step) -{ - window.move(step); -} - /** * @brief This class provides description of tile windowed view on the device memory. * @@ -1242,13 +1225,4 @@ make_tile_window_raw(const tile_window_with_static_lengths -CK_TILE_DEVICE void move_tile_window( - tile_window_with_static_lengths& window, - const typename tile_window_with_static_lengths::BottomTensorIndex& - step) -{ - window.move(step); -} - } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tile_window_linear.hpp b/include/ck_tile/core/tensor/tile_window_linear.hpp index 1e24e660f6..6af6813e0c 100644 --- a/include/ck_tile/core/tensor/tile_window_linear.hpp +++ b/include/ck_tile/core/tensor/tile_window_linear.hpp @@ -1200,19 +1200,4 @@ make_tile_window_linear_raw(const TileWindow_& tile_window, LinearBottomDims_{}); } -template -CK_TILE_DEVICE void move_tile_window( - tile_window_linear& - window, - const typename tile_window_linear::BottomTensorIndex& step) -{ - window.move(step); -} - } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tile_window_utils.hpp b/include/ck_tile/core/tensor/tile_window_utils.hpp index 71a72329f8..a6d4fcde36 100644 --- a/include/ck_tile/core/tensor/tile_window_utils.hpp +++ b/include/ck_tile/core/tensor/tile_window_utils.hpp @@ -18,6 +18,14 @@ #pragma once namespace ck_tile { +template +CK_TILE_DEVICE void move_tile_window( + TileWindow_& window, + const typename TileWindow_::BottomTensorIndex& step) +{ + window.move(step); +} + // input a lds store tile, extract some information from it // used to set m0 value for gfx9 serious template diff --git a/include/ck_tile/core/tensor/transpose_tile.hpp b/include/ck_tile/core/tensor/transpose_tile.hpp index 5b65b79c1a..f34efe5c2f 100644 --- a/include/ck_tile/core/tensor/transpose_tile.hpp +++ b/include/ck_tile/core/tensor/transpose_tile.hpp @@ -83,6 +83,9 @@ CK_TILE_DEVICE void transpose_tile2d_impl_in_thread(OutTensor& out_tensor, constexpr index_t num_vec_in = vec_length_out; constexpr index_t num_vec_out = vec_length_in; + using InVec = array; + using OutVec = array; + // SFC constexpr auto scalars_per_access_arr = generate_array( [&](auto i) { return (i == y_dim_vec_in or i == y_dim_vec_out) ? y_lengths[i] : 1; }, @@ -98,84 +101,51 @@ CK_TILE_DEVICE void transpose_tile2d_impl_in_thread(OutTensor& out_tensor, static_assert(num_access > 0, "wrong! num_access should be larger than 0"); - if constexpr(num_vec_in == 1 || num_vec_out == 1) - { - // loop over SFC - static_for<0, num_access, 1>{}([&](auto iAccess) { - // data index [y0, y1, ...] in the order of input tensor - constexpr auto idx_y = SFC_Y::get_index(iAccess); + // in/out vectors to be transposed + thread_buffer in_vectors; + thread_buffer out_vectors; - constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y); - constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y); + // loop over SFC and do transpose + static_for<0, num_access, 1>{}([&](auto iAccess) { + // data index [y0, y1, ...] in the order of input tensor + constexpr auto idx_y_start = SFC_Y::get_index(iAccess); - if constexpr(vec_length_in == 1) - { - out_tensor.get_thread_buffer()[number{}] = - in_tensor.get_thread_buffer()[number{}]; - } - else - { - using Vec = array; - out_tensor.get_thread_buffer().template get_as( - number{}) = - in_tensor.get_thread_buffer().template get_as( - number{}); - } + // get input vectors + static_for<0, num_vec_in, 1>{}([&](auto i) { + constexpr auto idx_y_in = generate_tuple( + [&](auto ii) { + return ii == y_dim_vec_out ? idx_y_start[ii] + i : idx_y_start[ii]; + }, + number{}); + + constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y_in); + static_assert(in_offset % vec_length_in == 0); + + in_vectors(i).template get_as()(I0) = + in_tensor.get_thread_buffer() + .template get_as()[number{}]; }); - } - else - { - using InVec = array; - using OutVec = array; - // in/out vectors to be transposed - thread_buffer in_vectors; - thread_buffer out_vectors; + // transpose + transpose_vectors{}(in_vectors, out_vectors); - // loop over SFC and do transpose - static_for<0, num_access, 1>{}([&](auto iAccess) { - // data index [y0, y1, ...] in the order of input tensor - constexpr auto idx_y_start = SFC_Y::get_index(iAccess); + // set output vectors + static_for<0, num_vec_out, 1>{}([&](auto i) { + constexpr auto idx_y_out_tmp = generate_array( + [&](auto ii) { return ii == y_dim_vec_in ? idx_y_start[ii] + i : idx_y_start[ii]; }, + number{}); - // get input vectors - static_for<0, num_vec_in, 1>{}([&](auto i) { - constexpr auto idx_y_in = generate_tuple( - [&](auto ii) { - return ii == y_dim_vec_out ? idx_y_start[ii] + i : idx_y_start[ii]; - }, - number{}); + constexpr auto idx_y_out = + container_reorder_given_new2old(idx_y_out_tmp, y_dim_out_to_in); - constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y_in); - static_assert(in_offset % vec_length_in == 0); + constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y_out); + static_assert(out_offset % vec_length_out == 0); - in_vectors(i).template get_as()(I0) = - in_tensor.get_thread_buffer() - .template get_as()[number{}]; - }); - - // transpose - transpose_vectors{}(in_vectors, out_vectors); - - // set output vectors - static_for<0, num_vec_out, 1>{}([&](auto i) { - constexpr auto idx_y_out_tmp = generate_array( - [&](auto ii) { - return ii == y_dim_vec_in ? idx_y_start[ii] + i : idx_y_start[ii]; - }, - number{}); - - constexpr auto idx_y_out = - container_reorder_given_new2old(idx_y_out_tmp, y_dim_out_to_in); - - constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y_out); - static_assert(out_offset % vec_length_out == 0); - - out_tensor.get_thread_buffer().template set_as( - number{}, - out_vectors[i].template get_as()[I0]); - }); + out_tensor.get_thread_buffer().template set_as( + number{}, + out_vectors[i].template get_as()[I0]); }); - } + }); } } // namespace detail diff --git a/include/ck_tile/host/check_err.hpp b/include/ck_tile/host/check_err.hpp index 745c18d6dd..3453743427 100644 --- a/include/ck_tile/host/check_err.hpp +++ b/include/ck_tile/host/check_err.hpp @@ -186,7 +186,7 @@ check_err(const Range& out, { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; @@ -246,7 +246,7 @@ check_err(const Range& out, { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; @@ -305,7 +305,7 @@ check_err(const Range& out, { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; @@ -360,7 +360,7 @@ std::enable_if_t<(std::is_same_v, ranges::range_val { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; @@ -437,7 +437,7 @@ std::enable_if_t<(std::is_same_v, ranges::range_val { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i << "] != ref[" << i << "]: " << o_fp64 << " != " << r_fp64 << std::endl; @@ -495,7 +495,7 @@ std::enable_if_t<(std::is_same_v, ranges::range_val { max_err = err > max_err ? err : max_err; err_count++; - if(err_count < 5) + if(err_count < 5000000) { std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; diff --git a/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp b/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp index 59ff88c4f2..bd59fcecc7 100644 --- a/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp +++ b/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp @@ -98,6 +98,7 @@ __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_, int row = idx / N; // Compute row index int col = idx % N; // Compute column index + (void)Num_tokens; // assert(p_sorted_expert_ids_ != nullptr); // assert(TopK == 1); // assert(Num_tokens == 128); diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 6ed0de57ab..bc5d50e921 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -119,6 +119,17 @@ struct CShuffleEpilogue return kMWave * kNWave * kMPerXdl * kNPerXdl * sizeof(ODataType); } + CK_TILE_HOST_DEVICE static constexpr auto GetCDisrtribution() + { + using TileEncodingPattern = + TileDistributionEncodingPattern2D; + return TileEncodingPattern::Make2DStaticTileDistribution(); + } + template {})>{} + 0; // printf("idx_y_start:%d \n", idx_m); - constexpr auto mIter = number{}) / (kMPerXdl * kMWave)>{}; + constexpr auto MPerAcess = kMPerXdl * kMWave; + constexpr auto mIter = number{}) / (MPerAcess)>{}; + using CDstrEncode = typename decltype(dram_tile_distribution)::DstrEncode; + constexpr ck_tile::index_t MRepeat = CDstrEncode::hs_lengthss_[number<0>{}][number<0>{}]; - statically_indexed_array offsets; - static_for<0, 2 /*CMrepeats*/, 1>{}([&](auto m0) { + statically_indexed_array offsets; + + static_for<0, MRepeat, 1>{}([&](auto m0) { auto token_id = token_pos + m0 + c_coord[0] + mIter * kMPerXdl * kMWave; + // auto token_id = token_pos + c_coord[0] + mIter * MPerAcess + MPerAcess / MRepeat * m0.value; auto fused_token = p_sorted_tokens_id[token_id]; - index_t token_offset = fused_token & 0xffffff; - + index_t scatter_token_id = fused_token & 0xffffff; if constexpr(IsInputGemm) { - token_offset = token_offset * 3 /*TopK*/ + (fused_token >> 24); + scatter_token_id = scatter_token_id * TopK + (fused_token >> 24); } - offsets[m0] = token_offset * 4096; // Problem::kN_; + offsets[m0] = scatter_token_id * stride_C; // Problem::kN_; }); // printf("c_coord[number<0>{}]: %d \n", coord[number<0>{}]); // printf("mIter: %d", mIter+0); @@ -212,7 +229,13 @@ struct CShuffleEpilogue if constexpr(out_memory_data_op == memory_operation_enum::set) { - store_tile(out_dram_window, c_out_tensor, offsets); + auto tile_window = make_tile_scatter_gather(out_dram_window.get_bottom_tensor_view(), + out_dram_window.get_window_lengths(), + out_dram_window.get_window_origin(), + dram_tile_distribution, + offsets); + tile_window.store(c_out_tensor); + // store_tile(out_dram_window, c_out_tensor, offsets); } else { diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index a66d4845a8..794f7f21f2 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -26,7 +26,6 @@ #include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp" -#include "ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp" @@ -40,7 +39,6 @@ #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm/pipeline//moe_gemm_pipeline_ag_bg_cr.hpp" #include "ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp" #include "ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm.hpp" diff --git a/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp deleted file mode 100644 index bff913cfb8..0000000000 --- a/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp +++ /dev/null @@ -1,506 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/core/numeric/math.hpp" -#include "ck_tile/core/utility/literals.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" -#include "ck_tile/host.hpp" - -namespace ck_tile { - -struct MoeGemmHostArgs : public ck_tile::GemmHostArgs -{ - ck_tile::index_t NumTokens; - ck_tile::index_t TopK; - const ck_tile::index_t* p_sorted_token_ids; - const ck_tile::index_t* p_sorted_expert_ids; - const ck_tile::index_t* p_max_token_id; - - // TODO: add kbatch for splitk - CK_TILE_HOST MoeGemmHostArgs() noexcept = default; - CK_TILE_HOST MoeGemmHostArgs(const ck_tile::index_t* p_sorted_token_ids_, - const ck_tile::index_t* p_sorted_expert_ids_, - const ck_tile::index_t* p_max_token_id_, - const void* a_ptr_, - const void* b_ptr_, - void* c_ptr_, - ck_tile::index_t NumTokens_, - ck_tile::index_t TopK_, - ck_tile::index_t M_, - ck_tile::index_t N_, - ck_tile::index_t K_, - ck_tile::index_t stride_A_, - ck_tile::index_t stride_B_, - ck_tile::index_t stride_C_) - : GemmHostArgs(a_ptr_, b_ptr_, c_ptr_, 1, M_, N_, K_, stride_A_, stride_B_, stride_C_), - NumTokens(NumTokens_), - TopK(TopK_), - p_sorted_token_ids(p_sorted_token_ids_), - p_sorted_expert_ids(p_sorted_expert_ids_), - p_max_token_id(p_max_token_id_) - { - } - - private: - static constexpr index_t KBatch = 1; -}; - -template -struct MoeGemmKernel : public GemmKernel -{ - using TilePartitioner = remove_cvref_t; - using GemmPipeline = remove_cvref_t; - using EpiloguePipeline = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; - - static constexpr bool IsInputGemm = IsInputGemm_; - - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - - using OffsetTile1DPartitioner = OffsettedTile1DPartitioner; - using Base = GemmKernel; - using GemmKernelArgs = typename Base::GemmKernelArgs; - - using SplitKBatchOffset = typename Base::SplitKBatchOffset; - - static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize; - - static constexpr auto I0 = number<0>(); - static constexpr auto I1 = number<1>(); - static constexpr auto I2 = number<2>(); - - struct MoeGemmKernelArgs : public GemmKernelArgs - { - const ck_tile::index_t* p_sorted_token_ids; - const ck_tile::index_t* p_sorted_expert_ids; - const ck_tile::index_t* p_max_token_id; - ck_tile::index_t NumTokens; - ck_tile::index_t TopK; - - CK_TILE_HOST MoeGemmKernelArgs() noexcept = default; - CK_TILE_HOST MoeGemmKernelArgs(const ck_tile::index_t* p_sorted_token_ids_, - const ck_tile::index_t* p_sorted_expert_ids_, - const ck_tile::index_t* p_max_token_id_, - const void* a_ptr_, - const void* b_ptr_, - void* c_ptr_, - ck_tile::index_t NumTokens_, - ck_tile::index_t TopK_, - ck_tile::index_t M_, - ck_tile::index_t N_, - ck_tile::index_t K_, - ck_tile::index_t stride_A_, - ck_tile::index_t stride_B_, - ck_tile::index_t stride_C_, - ck_tile::index_t KBatch) - : GemmKernelArgs{a_ptr_, - b_ptr_, - c_ptr_, - M_, - N_, - K_, - stride_A_, - stride_B_, - stride_C_, - KBatch}, - p_sorted_token_ids(p_sorted_token_ids_), - p_sorted_expert_ids(p_sorted_expert_ids_), - p_max_token_id(p_max_token_id_), - NumTokens(NumTokens_), - TopK(TopK_) - { - } - - CK_TILE_HOST static constexpr MoeGemmKernelArgs - MakeKernelArgs(const MoeGemmHostArgs& hostArgs) - { - printf("in moe gemm kernel args!"); - return MoeGemmKernelArgs{hostArgs.p_sorted_token_ids, - hostArgs.p_sorted_expert_ids, - hostArgs.p_max_token_id, - hostArgs.a_ptr, - hostArgs.b_ptr, - hostArgs.c_ptr, - hostArgs.NumTokens, - hostArgs.TopK, - hostArgs.M, - hostArgs.N, - hostArgs.K, - hostArgs.stride_A, - hostArgs.stride_B, - hostArgs.stride_C, - 1 - /*hostArgs.k_batch*/}; - } - }; - - [[nodiscard]] CK_TILE_HOST static const std::string GetName() - { - // clang-format off - using P_ = GemmPipeline; - return concat('_', "moe_gemm", gemm_prec_str, - concat('x', P_::MPerBlock, P_::NPerBlock, P_::KPerBlock), - concat('x', P_::GetVectorSizeA(), P_::GetVectorSizeB(), P_::GetVectorSizeC()), - concat('x', P_::kPadM, P_::kPadN, P_::kPadK)); - // clang-format on - } - - __host__ static constexpr auto BlockSize() -> dim3 { return dim3(KernelBlockSize); } - - __host__ static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) - { - // TODO: remove assertion - assert(KBatch == 1); - return Base::GridSize(M, N, KBatch); - } - - CK_TILE_HOST_DEVICE static constexpr auto GetSmemSize() -> index_t - { - return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); - } - - template - CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr, - const BDataType* b_ptr, - CDataType* c_ptr, - const MoeGemmKernelArgs& kargs, - const SplitKBatchOffset& splitk_batch_offset) - { - static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!"); - const auto& a_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - a_ptr, - make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK, - splitk_batch_offset.splitted_k), - make_tuple(kargs.stride_A, 1), - number{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - a_ptr, - make_tuple(splitk_batch_offset.splitted_k, - IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK), - make_tuple(kargs.stride_A, 1), - number{}, - number<1>{}); - } - }(); - - const auto& b_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - if constexpr(TilePartitioner::BlockGemmShape::PermuteB) - { - constexpr index_t K1 = GemmPipeline::GetSmemPackB(); - const index_t K0 = splitk_batch_offset.splitted_k / K1; - constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB()); - const auto b_k0_n_k1_desc = - make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), - make_tuple(kargs.N * K1, K1, I1), - number{}, - number<1>{}); - const auto b_n_k_desc = transform_tensor_descriptor( - b_k0_n_k1_desc, - make_tuple(make_merge_transform(make_tuple(K0, K1)), - make_pass_through_transform(kargs.N)), - make_tuple(sequence<0, 2>{}, sequence<1>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - return make_tensor_view(b_ptr, b_n_k_desc); - } - else - { - return make_naive_tensor_view( - b_ptr, - make_tuple(splitk_batch_offset.splitted_k, kargs.N), - make_tuple(kargs.stride_B, 1), - number{}, - number<1>{}); - } - } - else - { - if constexpr(TilePartitioner::BlockGemmShape::PermuteB) - { - constexpr index_t K1 = GemmPipeline::GetSmemPackB(); - const index_t K0 = splitk_batch_offset.splitted_k / K1; - constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB()); - const auto b_k0_n_k1_desc = - make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), - make_tuple(kargs.N * K1, K1, I1), - number{}, - number<1>{}); - const auto b_n_k_desc = transform_tensor_descriptor( - b_k0_n_k1_desc, - make_tuple(make_merge_transform(make_tuple(K0, K1)), - make_pass_through_transform(kargs.N)), - make_tuple(sequence<0, 2>{}, sequence<1>{}), - make_tuple(sequence<1>{}, sequence<0>{})); - return make_tensor_view(b_ptr, b_n_k_desc); - } - else - { - return make_naive_tensor_view( - b_ptr, - make_tuple(kargs.N, splitk_batch_offset.splitted_k), - make_tuple(kargs.stride_B, 1), - number{}, - number<1>{}); - } - } - }(); - - // TODO: enable vector write for C in ColMajor - const auto& c_tensor_view = [&]() { - if constexpr(std::is_same_v) - { - return make_naive_tensor_view( - c_ptr, - make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens, - kargs.N), - make_tuple(kargs.stride_C, 1), - number{}, - number<1>{}); - } - else - { - return make_naive_tensor_view( - c_ptr, - make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken, - kargs.N), - make_tuple(1, kargs.stride_C), - number<1>{}, - number<1>{}); - } - }(); - - return make_tuple(a_tensor_view, b_tensor_view, c_tensor_view); - } - - template - CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id) - { - if constexpr(std::is_same_v) - return transform_tensor_view( - view, - make_tuple(make_indexing_transform( - view.get_tensor_descriptor().get_length(number<0>()), token_id), - make_pass_through_transform( - view.get_tensor_descriptor().get_length(number<1>()))), - make_tuple(sequence<0>{}, sequence<1>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - else - return transform_tensor_view( - view, - make_tuple(make_pass_through_transform( - view.get_tensor_descriptor().get_length(number<0>())), - make_indexing_transform( - view.get_tensor_descriptor().get_length(number<1>()), token_id)), - make_tuple(sequence<0>{}, sequence<1>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - } - - template - CK_TILE_DEVICE static auto GetCTransformGemmView(const CView& view, const index_t token_id) - { - if constexpr(std::is_same_v) - return transform_tensor_view( - view, - make_tuple(make_indexing_transform( - view.get_tensor_descriptor().get_length(number<0>()), token_id), - make_pass_through_transform( - view.get_tensor_descriptor().get_length(number<1>()))), - make_tuple(sequence<0>{}, sequence<1>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - else - return transform_tensor_view( - view, - make_tuple(make_pass_through_transform( - view.get_tensor_descriptor().get_length(number<0>())), - make_indexing_transform( - view.get_tensor_descriptor().get_length(number<1>()), token_id)), - make_tuple(sequence<0>{}, sequence<1>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - } - - template - CK_TILE_DEVICE static auto TransformGemmPadViews(const PadView& views, const index_t token_id) - { - auto a_pad_view = views.at(number<0>()); - auto b_pad_view = views.at(number<1>()); - auto c_pad_view = views.at(number<2>()); - - const auto a_gather_view = GetATransformGemmView(a_pad_view, token_id); - // TODO： Caculate expert offset of the buf in B. - - // const auto c_scatter_view = GetCTransformGemmView(c_pad_view, token_id); - // if (token_id){} - return make_tuple(a_gather_view, b_pad_view, c_pad_view); - } - - template - CK_TILE_DEVICE static auto - MakeGemmTileWindows(const PadView& views, const index_t i_m, const index_t i_n) - { - const auto& a_pad_view = views.at(number<0>{}); - const auto& b_pad_view = views.at(number<1>{}); - const auto& c_pad_view = views.at(number<2>{}); - if(i_m) {} - const auto& a_block_window = [&]() { - if constexpr(std::is_same_v) - { - return make_tile_window(a_pad_view, - make_tuple(number{}, - number{}), - {0, 0}); - } - else - { - return make_tile_window(a_pad_view, - make_tuple(number{}, - number{}), - {0, 0}); - } - }(); - - const auto& b_block_window = [&]() { - if constexpr(std::is_same_v) - { - return make_tile_window(b_pad_view, - make_tuple(number{}, - number{}), - {i_n, 0}); - } - else - { - return make_tile_window(b_pad_view, - make_tuple(number{}, - number{}), - {0, i_n}); - } - }(); - - auto c_block_window = make_tile_window( - c_pad_view, - make_tuple(number{}, number{}), - {0, i_n}); - - return make_tuple(a_block_window, b_block_window, c_block_window); - } - - template - CK_TILE_DEVICE void operator()(const MoeGemmKernelArgs gemm_desc) const - { - // TODO: implement C scatter store accordring to expert_id - // TODO: the branch without swizzle - const index_t max_token_id = __builtin_amdgcn_readfirstlane(gemm_desc.p_max_token_id[0]); - const index_t block_id = ck_tile::get_block_1d_id(); - - // TODO: check the block id caculation - const auto [expert_blk_id, _] = - OffsetTile1DPartitioner::GetOffsetedTileIndex(0, gemm_desc.M, gemm_desc.N); - - if(expert_blk_id * TilePartitioner::MPerBlock >= max_token_id) - return; - - const index_t NBlocks = gemm_desc.N / TilePartitioner::NPerBlock; - const index_t expert_id = gemm_desc.p_sorted_expert_ids[expert_blk_id]; - const index_t prefix_blk_m = gemm_desc.p_max_token_id[1 + expert_id]; - const index_t blk_cnt_of_eid = gemm_desc.p_max_token_id[2 + expert_id]; - - // printf("expert_blk_id: %d, expert_id: %d \n",expert_blk_id, expert_id); - - // expert_id = expert_blk_id; - - const index_t block_start = prefix_blk_m * NBlocks; - - const index_t ecnt = blk_cnt_of_eid - prefix_blk_m; - const index_t expert_swizzle = ecnt > 0 ? ecnt : 1; - // index_t block_end = block_start + blk_cnt_of_eid * NBlocks; - - const index_t block_id_start_in_expert = block_id - block_start; - const index_t im = __builtin_amdgcn_readfirstlane(prefix_blk_m + block_id_start_in_expert / - 8 % expert_swizzle); - const index_t in = __builtin_amdgcn_readfirstlane( - block_id_start_in_expert % 8 + block_id_start_in_expert / (8 * expert_swizzle) * 8); - - const auto a_coord = GemmPipeline::GetACoord(); // 2d thread offset, [i_row, i_col] - const auto sorted_token_id = a_coord[number<0>{}] + im * TilePartitioner::MPerBlock; - - // constexpr auto AMRepeat = GemmPipeline::GetAMRepeat(); - - // ck_tile::statically_indexed_array gather_offset; - // static_for<0, AMRepeat, 1>{}([&](auto thr_offset_m){ - // const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id + - // thr_offset_m]; gather_offset(thr_offset_m) = fused_token & 0xffffff; - // }); - - const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id]; - - // TODO: token_id should include topk offset depends on ffn1 or ffn2 - const index_t token_id = fused_token & 0xffffff; - - if constexpr(!IsInputGemm) - { - token_id = token_id * gemm_desc.TopK + (fused_token >> 24); - } - - const index_t expert_stride = __builtin_amdgcn_readfirstlane(gemm_desc.N * gemm_desc.K); - - const typename Base::SplitKBatchOffset splitk_batch_offset(gemm_desc); - // options - const ADataType* a_ptr = - static_cast(gemm_desc.a_ptr) + splitk_batch_offset.a_k_split_offset; - const BDataType* b_ptr = static_cast(gemm_desc.b_ptr) + - splitk_batch_offset.b_k_split_offset + expert_stride * expert_id; - CDataType* c_ptr = static_cast(gemm_desc.c_ptr); - - const auto& gemm_tensor_views_tuple = - MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, gemm_desc, splitk_batch_offset); - const auto& gemm_pad_views = Base::MakeGemmPadViews(gemm_tensor_views_tuple); - const auto& transformed_views = TransformGemmPadViews(gemm_pad_views, token_id); - auto gemm_tile_windows = MakeGemmTileWindows( - transformed_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock); - const index_t num_loop = - __builtin_amdgcn_readfirstlane(TilePartitioner::GetLoopNum(gemm_desc.K)); - - // printf("num_loop: %d", num_loop); - - static_assert(GemmPipeline::DoubleSmemBuffer == true, - "For now, only support doublesmembuffer"); - - __shared__ char smem_ptr_0[GetSmemSize()]; - __shared__ char smem_ptr_1[GetSmemSize()]; - // Run GEMM cooperatively by whole workgroup. - const auto& a_block_window = gemm_tile_windows.at(number<0>{}); - const auto& b_block_window = gemm_tile_windows.at(number<1>{}); - - const auto& c_block_tile = GemmPipeline{}.template operator()( - a_block_window, b_block_window, num_loop, smem_ptr_0, smem_ptr_1); - - // Run Epilogue Pipeline - auto& c_block_window = gemm_tile_windows.at(number<2>{}); - - EpiloguePipeline{}.template operator()( - c_block_window, - c_block_tile, - smem_ptr_0, - gemm_desc.p_sorted_token_ids, - im * TilePartitioner::MPerBlock); - } -}; - -} // namespace ck_tile diff --git a/include/ck_tile/ops/moe_gemm.hpp b/include/ck_tile/ops/moe_gemm.hpp new file mode 100644 index 0000000000..7819cb412e --- /dev/null +++ b/include/ck_tile/ops/moe_gemm.hpp @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once +#include "ck_tile/ops/flatmm/block/block_flatmm_asmem_bsmem_creg_v1.hpp" +#include "ck_tile/ops/flatmm/block/block_flatmm_asmem_bsmem_creg_v1_custom_policy.hpp" +#include "ck_tile/ops/flatmm/pipeline/tile_flatmm_shape.hpp" +#include "ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp" + +#include "ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp" +#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp" +#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp" +#include "ck_tile/ops/common/generic_2d_block_shape.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/common/utils.hpp" diff --git a/include/ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp b/include/ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp new file mode 100644 index 0000000000..622b323839 --- /dev/null +++ b/include/ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp @@ -0,0 +1,649 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core/numeric/math.hpp" +#include "ck_tile/core/utility/literals.hpp" +#include "ck_tile/ops/flatmm/kernel/flatmm_kernel.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp" +#include "ck_tile/host.hpp" + +// #define disable_tile_gs + +namespace ck_tile { + +struct MoeGemmHostArgs : public ck_tile::FlatmmHostArgs +{ + ck_tile::index_t NumTokens; + ck_tile::index_t TopK; + const ck_tile::index_t* p_sorted_token_ids; + const ck_tile::index_t* p_sorted_expert_ids; + const ck_tile::index_t* p_max_token_id; + + // TODO: add kbatch for splitk + CK_TILE_HOST MoeGemmHostArgs() noexcept = default; + CK_TILE_HOST MoeGemmHostArgs(const ck_tile::index_t* p_sorted_token_ids_, + const ck_tile::index_t* p_sorted_expert_ids_, + const ck_tile::index_t* p_max_token_id_, + const void* a_ptr_, + const void* b_shuffle_ptr_, + void* c_ptr_, + ck_tile::index_t NumTokens_, + ck_tile::index_t TopK_, + ck_tile::index_t k_batch_, + ck_tile::index_t M_, + ck_tile::index_t N_, + ck_tile::index_t K_, + ck_tile::index_t stride_A_, + ck_tile::index_t stride_B_, + ck_tile::index_t stride_C_) + : FlatmmHostArgs(a_ptr_, b_shuffle_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_), + NumTokens(NumTokens_), + TopK(TopK_), + p_sorted_token_ids(p_sorted_token_ids_), + p_sorted_expert_ids(p_sorted_expert_ids_), + p_max_token_id(p_max_token_id_) + { + } + // TODO: why kBatch? + // private: + // static constexpr index_t KBatch = 1; +}; + +template +struct MoeGemmKernel +{ + using TilePartitioner = remove_cvref_t; + using FlatmmPipeline = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + using ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + using BlockGemmShape = + remove_cvref_t; // TileFlatmmShape + + static constexpr bool IsInputGemm = IsInputGemm_; + + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using OffsetTile1DPartitioner = OffsettedTile1DPartitioner; + static constexpr index_t KernelBlockSize = FlatmmPipeline::BlockSize; + + static constexpr auto I0 = number<0>(); + static constexpr auto I1 = number<1>(); + static constexpr auto I2 = number<2>(); + + struct MoeGemmKernelArgs + { + const ck_tile::index_t* p_sorted_token_ids; + const ck_tile::index_t* p_sorted_expert_ids; + const ck_tile::index_t* p_max_token_id; + const void* p_a_ptr; + const void* p_b_shuffle_ptr; + void* p_c_ptr; + ck_tile::index_t NumTokens; + ck_tile::index_t TopK; + ck_tile::index_t M; + ck_tile::index_t N; + ck_tile::index_t K; + ck_tile::index_t stride_A; + ck_tile::index_t stride_B; + ck_tile::index_t stride_C; + ck_tile::index_t k_batch; + // + // CK_TILE_HOST MoeGemmKernelArgs() noexcept = default; + // CK_TILE_HOST MoeGemmKernelArgs(const ck_tile::index_t* p_sorted_token_ids_, + // const ck_tile::index_t* p_sorted_expert_ids_, + // const ck_tile::index_t* p_max_token_id_, + // const void* a_ptr_, + // const void* b_shuffle_ptr_, + // void* c_ptr_, + // ck_tile::index_t NumTokens_, + // ck_tile::index_t TopK_, + // ck_tile::index_t M_, + // ck_tile::index_t N_, + // ck_tile::index_t K_, + // ck_tile::index_t stride_A_, + // ck_tile::index_t stride_B_, + // ck_tile::index_t stride_C_, + // ck_tile::index_t KBatch) : + // p_sorted_token_ids(p_sorted_token_ids_), + // p_sorted_expert_ids(p_sorted_expert_ids_), + // p_max_token_id(p_max_token_id_), + // p_a_ptr(a_ptr_), + // p_b_shuffle_ptr(b_shuffle_ptr_), + // p_c_ptr(c_ptr_), + // NumTokens(NumTokens_), + // TopK(TopK_), + // M(M_), + // N(N_), + // K(K_), + // stride_A(stride_A_), + // stride_B(stride_B_), + // stride_C(stride_C_), + // k_batch(KBatch) + // { + // } + + }; + + CK_TILE_HOST static constexpr MoeGemmKernelArgs MakeKernelArgs(const MoeGemmHostArgs& hostArgs) + { + printf("in moe gemm kernel args! \n"); + return MoeGemmKernelArgs{hostArgs.p_sorted_token_ids, + hostArgs.p_sorted_expert_ids, + hostArgs.p_max_token_id, + hostArgs.a_ptr, + hostArgs.b_shuffle_ptr, + hostArgs.c_ptr, + hostArgs.NumTokens, + hostArgs.TopK, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_A, + hostArgs.stride_B, + hostArgs.stride_C, + 1 + /*hostArgs.k_batch*/}; + } + + [[nodiscard]] CK_TILE_HOST static const std::string GetName() + { + // clang-format off + using P_ = FlatmmPipeline; + return concat('_', "moe_gemm", gemm_prec_str, + concat('x', P_::kMPerBlock, P_::kNPerBlock, P_::kKPerBlock), + concat('x', P_::GetVectorSizeA(), P_::GetVectorSizeB(), P_::GetVectorSizeC()), + concat('x', P_::kPadM, P_::kPadN, P_::kPadK)); + // clang-format on + } + + __host__ static constexpr auto BlockSize() -> dim3 { return dim3(KernelBlockSize); } + + __host__ static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) + { + // TODO: remove assertion + assert(KBatch == 1); + return dim3(TilePartitioner::GridSize(M, N), 1, KBatch); + } + + CK_TILE_HOST_DEVICE static constexpr auto GetSmemSize() -> index_t + { + return max(FlatmmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); + } + + struct SplitKBatchOffset + { + __device__ SplitKBatchOffset(const MoeGemmKernelArgs& kargs, + const std::size_t k_id = blockIdx.z) + { + constexpr auto K1 = TilePartitioner::BlockGemmShape::WarpTile::at(number<2>{}); + const index_t K_t = __builtin_amdgcn_readfirstlane(kargs.k_batch * K1); + const index_t KRead = __builtin_amdgcn_readfirstlane((kargs.K + K_t - 1) / K_t * K1); + + if constexpr(std::is_same_v) + { + a_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead); + } + else if constexpr(std::is_same_v) + { + a_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead * kargs.stride_A); + } + + if constexpr(std::is_same_v) + { + b_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead * kargs.stride_B); + } + else if constexpr(std::is_same_v) + { + b_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead); + } + + if(k_id < static_cast(kargs.k_batch - 1)) + { + splitted_k = __builtin_amdgcn_readfirstlane(KRead); + } + else + { + splitted_k = __builtin_amdgcn_readfirstlane(kargs.K - KRead * (kargs.k_batch - 1)); + } + } + + index_t a_k_split_offset; + index_t b_k_split_offset; + index_t splitted_k; + }; + + template + CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr, + const BDataType* b_flat_ptr, + CDataType* c_ptr, + const MoeGemmKernelArgs& kargs, + const SplitKBatchOffset& splitk_batch_offset) + { + // static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!"); + const auto& a_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + a_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK, + splitk_batch_offset.splitted_k), + make_tuple(kargs.stride_A, 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + a_ptr, + make_tuple(splitk_batch_offset.splitted_k, + IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK), + make_tuple(kargs.stride_A, 1), + number{}, + number<1>{}); + } + }(); + + index_t kFlatK = FlatmmPipeline::flatKPerWarp * (splitk_batch_offset.splitted_k / + BlockGemmShape::WarpTile::at(number<2>{})); + index_t kFlatN = kargs.N * kargs.K / kFlatK; + const auto& b_flat_tensor_view = [&]() { + return make_naive_tensor_view( + b_flat_ptr, + make_tuple(kFlatN, kFlatK), + make_tuple(kFlatK, 1), + number{}, + number<1>{}); + }(); + + + // const auto& b_tensor_view = [&]() { + // if constexpr(std::is_same_v) + // { + // if constexpr(TilePartitioner::BlockGemmShape::PermuteB) + // { + // constexpr index_t K1 = FlatmmPipeline::GetSmemPackB(); + // const index_t K0 = splitk_batch_offset.splitted_k / K1; + // constexpr index_t VectorSizeB = std::min(K1, FlatmmPipeline::GetVectorSizeB()); + // const auto b_k0_n_k1_desc = + // make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), + // make_tuple(kargs.N * K1, K1, I1), + // number{}, + // number<1>{}); + // const auto b_n_k_desc = transform_tensor_descriptor( + // b_k0_n_k1_desc, + // make_tuple(make_merge_transform(make_tuple(K0, K1)), + // make_pass_through_transform(kargs.N)), + // make_tuple(sequence<0, 2>{}, sequence<1>{}), + // make_tuple(sequence<0>{}, sequence<1>{})); + // return make_tensor_view(b_ptr, b_n_k_desc); + // } + // else + // { + // return make_naive_tensor_view( + // b_ptr, + // make_tuple(splitk_batch_offset.splitted_k, kargs.N), + // make_tuple(kargs.stride_B, 1), + // number{}, + // number<1>{}); + // } + // } + // else + // { + // if constexpr(TilePartitioner::BlockGemmShape::PermuteB) + // { + // constexpr index_t K1 = FlatmmPipeline::GetSmemPackB(); + // const index_t K0 = splitk_batch_offset.splitted_k / K1; + // constexpr index_t VectorSizeB = std::min(K1, FlatmmPipeline::GetVectorSizeB()); + // const auto b_k0_n_k1_desc = + // make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), + // make_tuple(kargs.N * K1, K1, I1), + // number{}, + // number<1>{}); + // const auto b_n_k_desc = transform_tensor_descriptor( + // b_k0_n_k1_desc, + // make_tuple(make_merge_transform(make_tuple(K0, K1)), + // make_pass_through_transform(kargs.N)), + // make_tuple(sequence<0, 2>{}, sequence<1>{}), + // make_tuple(sequence<1>{}, sequence<0>{})); + // return make_tensor_view(b_ptr, b_n_k_desc); + // } + // else + // { + // return make_naive_tensor_view( + // b_ptr, + // make_tuple(kargs.N, splitk_batch_offset.splitted_k), + // make_tuple(kargs.stride_B, 1), + // number{}, + // number<1>{}); + // } + // } + // }(); + + // TODO: enable vector write for C in ColMajor + const auto& c_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + c_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens, + kargs.N), + make_tuple(kargs.stride_C, 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + c_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken, + kargs.N), + make_tuple(1, kargs.stride_C), + number<1>{}, + number<1>{}); + } + }(); + + return make_tuple(a_tensor_view, b_flat_tensor_view, c_tensor_view); + } + + template + CK_TILE_DEVICE static auto MakeGemmPadViews(const TensorView& views) + { + const auto& a_pad_view = [&]() { + const auto& a_tensor_view = views.at(I0); + if constexpr(std::is_same_v) + { + return pad_tensor_view(a_tensor_view, + make_tuple(number{}, + number{}), + sequence{}); + } + else + { + return pad_tensor_view(a_tensor_view, + make_tuple(number{}, + number{}), + sequence{}); + } + }(); + + const auto& b_flat_tensor_view = views.at(I1); + + // TODO vector write in for C in ColMajor + const auto& c_pad_view = [&]() { + const auto& c_tensor_view = views.at(I2); + if constexpr(std::is_same_v) + { + return pad_tensor_view(c_tensor_view, + make_tuple(number{}, + number{}), + sequence{}); + } + else + { + return pad_tensor_view(c_tensor_view, + make_tuple(number{}, + number{}), + sequence{}); + } + }(); + + return make_tuple(a_pad_view, b_flat_tensor_view, c_pad_view); + } + + template + CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id) + { + if constexpr(std::is_same_v) + return transform_tensor_view( + view, + make_tuple(make_indexing_transform( + view.get_tensor_descriptor().get_length(number<0>()), token_id), + make_pass_through_transform( + view.get_tensor_descriptor().get_length(number<1>()))), + make_tuple(sequence<0>{}, sequence<1>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + else + return transform_tensor_view( + view, + make_tuple(make_pass_through_transform( + view.get_tensor_descriptor().get_length(number<0>())), + make_indexing_transform( + view.get_tensor_descriptor().get_length(number<1>()), token_id)), + make_tuple(sequence<0>{}, sequence<1>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + } + + // template + // CK_TILE_DEVICE static auto GetCTransformGemmView(const CView& view, const index_t token_id) + // { + // if constexpr(std::is_same_v) + // return transform_tensor_view( + // view, + // make_tuple(make_indexing_transform( + // view.get_tensor_descriptor().get_length(number<0>()), token_id), + // make_pass_through_transform( + // view.get_tensor_descriptor().get_length(number<1>()))), + // make_tuple(sequence<0>{}, sequence<1>{}), + // make_tuple(sequence<0>{}, sequence<1>{})); + // else + // return transform_tensor_view( + // view, + // make_tuple(make_pass_through_transform( + // view.get_tensor_descriptor().get_length(number<0>())), + // make_indexing_transform( + // view.get_tensor_descriptor().get_length(number<1>()), token_id)), + // make_tuple(sequence<0>{}, sequence<1>{}), + // make_tuple(sequence<0>{}, sequence<1>{})); + // } + + template + CK_TILE_DEVICE static auto TransformGemmPadViews(const PadView& views, const index_t token_id) + { + auto a_pad_view = views.at(number<0>()); + auto b_pad_view = views.at(number<1>()); + auto c_pad_view = views.at(number<2>()); + + const auto a_gather_view = GetATransformGemmView(a_pad_view, token_id); + // TODO： Caculate expert offset of the buf in B. + + // const auto c_scatter_view = GetCTransformGemmView(c_pad_view, token_id); + // if (token_id){} + return make_tuple(a_gather_view, b_pad_view, c_pad_view); + } + + template + CK_TILE_DEVICE static auto + MakeGemmTileWindows(const PadView& views, const index_t i_m, const index_t i_n) + { + (void)i_m; + + const auto& a_pad_view = views.at(number<0>{}); + const auto& b_flat_pad_view = views.at(number<1>{}); + const auto& c_pad_view = views.at(number<2>{}); + // if(i_m) {} + + const auto& a_block_window = [&]() { + if constexpr(std::is_same_v) + { + return make_tile_window(a_pad_view, + make_tuple(number{}, + number{}), + {0, 0}); + } + else + { + return make_tile_window(a_pad_view, + make_tuple(number{}, + number{}), + {0, 0}); + } + }(); + + const auto& b_flat_block_window = + make_tile_window(b_flat_pad_view, + make_tuple(number{}, + number{}), + {static_cast(i_n / BlockGemmShape::WarpTile::at(I1)), 0}); + + auto c_block_window = make_tile_window( + c_pad_view, + make_tuple(number{}, number{}), + // {i_m, i_n}); + {0, i_n}); + + return make_tuple(a_block_window, b_flat_block_window, c_block_window); + } + + template + CK_TILE_DEVICE void operator()(const MoeGemmKernelArgs gemm_desc) const + { + // TODO: implement C scatter store accordring to expert_id + // TODO: the branch without swizzle + const index_t max_token_id = __builtin_amdgcn_readfirstlane(gemm_desc.p_max_token_id[0]); + const index_t block_id = ck_tile::get_block_1d_id(); + + // TODO: check the block id caculation + const auto [expert_blk_id, _] = + OffsetTile1DPartitioner::GetOffsetedTileIndex(0, gemm_desc.M, gemm_desc.N); + + if(expert_blk_id * TilePartitioner::MPerBlock >= max_token_id) + return; + + const index_t NBlocks = gemm_desc.N / TilePartitioner::NPerBlock; + const index_t expert_id = gemm_desc.p_sorted_expert_ids[expert_blk_id]; + const index_t prefix_blk_m = gemm_desc.p_max_token_id[1 + expert_id]; + const index_t blk_cnt_of_eid = gemm_desc.p_max_token_id[2 + expert_id]; + + // printf("expert_blk_id: %d, expert_id: %d \n",expert_blk_id, expert_id); + + // expert_id = expert_blk_id; + + const index_t block_start = prefix_blk_m * NBlocks; + + const index_t ecnt = blk_cnt_of_eid - prefix_blk_m; + const index_t expert_swizzle = ecnt > 0 ? ecnt : 1; + // index_t block_end = block_start + blk_cnt_of_eid * NBlocks; + + const index_t block_id_start_in_expert = block_id - block_start; + const index_t im = __builtin_amdgcn_readfirstlane(prefix_blk_m + block_id_start_in_expert / + 8 % expert_swizzle); + const index_t in = __builtin_amdgcn_readfirstlane( + block_id_start_in_expert % 8 + block_id_start_in_expert / (8 * expert_swizzle) * 8); + + const auto a_coord = FlatmmPipeline::GetACoord(); // 2d thread offset, [i_row, i_col] +#ifdef disable_tile_gs + const auto sorted_token_id = a_coord[number<0>{}] + im * TilePartitioner::MPerBlock; + const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id]; + + // TODO: token_id should include topk offset depends on ffn1 or ffn2 + constexpr index_t token_id_mask = 0xffffff; + index_t token_id = fused_token & token_id_mask; + if constexpr(!IsInputGemm) + { + constexpr index_t token_id_offset = 24; + token_id = token_id * gemm_desc.TopK + (fused_token >> token_id_offset); + } +#else + constexpr ck_tile::index_t MRepeat = FlatmmPipeline::GetAMRepeat(); + statically_indexed_array a_offsets; + + constexpr index_t token_id_mask = 0xffffff; + constexpr index_t token_id_offset = 24; + + // constexpr auto kMWave = TilePartitioner::BlockGemmShape::BlockWarps::at(I0); + // constexpr auto kNWave = TilePartitioner::BlockGemmShape::BlockWarps::at(I1); + // const index_t iMWarp = get_warp_id() / kNWave; + static_for<0, MRepeat, 1>{}([&](auto m0) { + // const auto sorted_token_id = a_coord[I0] + im * TilePartitioner::MPerBlock + + // iMWarp * TilePartitioner::MPerBlock / kMWave + + // m0 * TilePartitioner::MPerBlock / kMWave / MRepeat; + const auto sorted_token_id = a_coord[I0] + im * TilePartitioner::MPerBlock + + m0 * TilePartitioner::MPerBlock / MRepeat; + const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id]; + + // TODO: token_id should include topk offset depends on ffn1 or ffn2 + index_t gather_token_id = fused_token & token_id_mask; + if constexpr(!IsInputGemm) + { + gather_token_id = gather_token_id * gemm_desc.TopK + (fused_token >> token_id_offset); + } + a_offsets[m0] = gather_token_id * gemm_desc.stride_A; + }); +#endif + + const index_t expert_stride = __builtin_amdgcn_readfirstlane(gemm_desc.N * gemm_desc.K); + + const SplitKBatchOffset splitk_batch_offset(gemm_desc); + // options + const ADataType* a_ptr = + static_cast(gemm_desc.p_a_ptr) + splitk_batch_offset.a_k_split_offset; + const BDataType* b_shuffle_ptr = static_cast(gemm_desc.p_b_shuffle_ptr) + + splitk_batch_offset.b_k_split_offset + expert_stride * expert_id; + CDataType* c_ptr = static_cast(gemm_desc.p_c_ptr); + + const auto& gemm_tensor_views_tuple = + MakeGemmTensorViews(a_ptr, b_shuffle_ptr, c_ptr, gemm_desc, splitk_batch_offset); + const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple); + +#ifdef disable_tile_gs + const auto& transformed_views = TransformGemmPadViews(gemm_pad_views, token_id); + auto gemm_tile_windows = MakeGemmTileWindows( + transformed_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock); +#else + auto gemm_tile_windows = MakeGemmTileWindows( + gemm_pad_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock); +#endif + + const index_t num_loop = + __builtin_amdgcn_readfirstlane(TilePartitioner::GetLoopNum(gemm_desc.K)); + + // printf("num_loop: %d", num_loop); + + // static_assert(FlatmmPipeline::DoubleSmemBuffer == true, + // "For now, only support doublesmembuffer"); + + __shared__ char smem_ptr_0[GetSmemSize()]; + // __shared__ char smem_ptr_1[GetSmemSize()]; + // Run GEMM cooperatively by whole workgroup. + const auto& a_block_window = gemm_tile_windows.at(number<0>{}); + const auto& b_block_window = gemm_tile_windows.at(number<1>{}); + +#ifdef disable_tile_gs + const auto& c_block_tile = FlatmmPipeline{}.template operator()( + a_block_window, b_block_window, num_loop, smem_ptr_0); +#else + auto a_gather_block_tile = ck_tile::make_tile_scatter_gather( + a_block_window.get_bottom_tensor_view(), + a_block_window.get_window_lengths(), + a_block_window.get_window_origin(), + FlatmmPipeline::GetADramTileDistribution(), + a_offsets); // K DRAM tile window for + const auto& c_block_tile = FlatmmPipeline{}.template operator()( + a_gather_block_tile, b_block_window, num_loop, smem_ptr_0); +#endif + + // Run Epilogue Pipeline + auto& c_block_window = gemm_tile_windows.at(number<2>{}); + + EpiloguePipeline{}.template operator()( + c_block_window, + c_block_tile, + smem_ptr_0, + gemm_desc.p_sorted_token_ids, + im * TilePartitioner::MPerBlock, + gemm_desc.TopK, + gemm_desc.stride_C); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp similarity index 100% rename from include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp rename to include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp diff --git a/include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp similarity index 88% rename from include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp rename to include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp index a587c9c083..58e6d7405b 100644 --- a/include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp @@ -6,10 +6,11 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp" namespace ck_tile { template -struct MoeGemmPipelineAgBgCrPolicy : public GemmPipelineAgBgCrCompV4DefaultPolicy +struct MoeGemmPipelineAgBgCrPolicy : public UniversalFlatmmPipelineAgBgCrPolicy { template CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_C() @@ -33,4 +34,4 @@ struct MoeGemmPipelineAgBgCrPolicy : public GemmPipelineAgBgCrCompV4DefaultPolic return c_block_dstr; } } -} // namespace ck_tile \ No newline at end of file +} // namespace ck_tile diff --git a/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp new file mode 100644 index 0000000000..6b807862e0 --- /dev/null +++ b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp @@ -0,0 +1,245 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common.hpp" +#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp" +#include + +namespace ck_tile { + +template +struct MoeGemmPipelineAgBgCrImpl +{ + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static_assert(!std::is_same_v, "Not implemented"); + + static constexpr index_t APackedSize = + ck_tile::numeric_traits>::PackedSize; + static constexpr index_t BPackedSize = + ck_tile::numeric_traits>::PackedSize; + + using ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using BlockFlatmm = remove_cvref_t())>; + using I0 = number<0>; + using I1 = number<1>; + using I2 = number<2>; + + static constexpr index_t BlockSize = Problem::kBlockSize; + + static constexpr index_t kMPerBlock = BlockGemmShape::kM; + static constexpr index_t kNPerBlock = BlockGemmShape::kN; + static constexpr index_t kKPerBlock = BlockGemmShape::kK; + + static constexpr index_t flatKPerWarp = BlockGemmShape::flatKPerWarp; + static constexpr index_t flatNPerWarp = BlockGemmShape::flatNPerWarp; + + // static constexpr index_t GetVectorSizeA() { return PipelinePolicy::template GetVectorSizeA(); } + // static constexpr index_t GetVectorSizeB() { return PipelinePolicy::template GetVectorSizeB(); } + // static constexpr index_t GetVectorSizeC() { return PipelinePolicy::template GetVectorSizeC(); } + + static constexpr index_t GetVectorSizeA() { return Problem::VectorSizeA; } + static constexpr index_t GetVectorSizeB() { return Problem::VectorSizeB; } + static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; } + static constexpr index_t GetSmemPackA() { return PipelinePolicy::template GetSmemPackA(); } + static constexpr index_t GetSmemPackB() { return PipelinePolicy::template GetSmemPackB(); } + + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool kPadK = Problem::kPadK; + + // static constexpr bool DoubleSmemBuffer = Problem::DoubleSmemBuffer; + + static constexpr bool HasHotLoop = Problem::HasHotLoop; + static constexpr auto Scheduler = Problem::Scheduler; + + CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; } + + CK_TILE_HOST_DEVICE static constexpr auto GetADramTileDistribution() { + return PipelinePolicy::template MakeADramTileDistribution(); + } + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return PipelinePolicy::template GetSmemSize(); + } + + CK_TILE_HOST_DEVICE constexpr static auto GetACoord() + { + constexpr auto a_dist = PipelinePolicy::template MakeADramTileDistribution(); + return a_dist.calculate_index(); + } + + // get thread coordinate of A in the threadblock + CK_TILE_HOST_DEVICE constexpr static auto GetAMRepeat() + { + constexpr auto a_dist = PipelinePolicy::template MakeADramTileDistribution(); + + using ADstrEncode = typename decltype(a_dist)::DstrEncode; + constexpr ck_tile::index_t MRepeat = ADstrEncode::hs_lengthss_[number<0>{}][number<0>{}]; + return MRepeat; + } + + template + CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindow& a_dram_block_window, + const AElementFunction& a_element_func, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + static_assert( + std::is_same_v>, + "wrong!"); + + static_assert(kMPerBlock == ADramBlockWindow{}.get_window_lengths()[number<0>{}], + "wrong!"); + static_assert(kKPerBlock == ADramBlockWindow{}.get_window_lengths()[number<1>{}], + "wrong!"); + + // A tile in LDS + ADataType* p_a_lds = static_cast(p_smem); + + constexpr auto a_lds_block_desc = + PipelinePolicy::template MakeALdsBlockDescriptor(); + + auto a_lds_block = make_tensor_view(p_a_lds, a_lds_block_desc); + + // auto a_dist = PipelinePolicy::template MakeADramTileDistribution(); + // auto a_coord = a_dist.calculate_index(); + // using ADstrEncode = typename decltype(a_dist)::DstrEncode; + // constexpr ck_tile::index_t MRepeat = ADstrEncode::hs_lengthss_[I0][I0]; + // statically_indexed_array a_offsets; + // static_for<0, MRepeat, 1>{}([&](auto n0) { + // int32_t seqlen_k_idx_per_repeat = cur_seqlen_k_idx + k_coord[0] + Traits::kBlockN / NRepeat * n0.value; + // int32_t page_idx = seqlen_k_idx_per_repeat / page_block_size; + // int32_t seq_idx = seqlen_k_idx_per_repeat % page_block_size; + // k_offsets[n0] = (block_indices[page_idx] * page_block_size + seq_idx) * stride_s_k; + // }); + // + // // A DRAM tile window for load + // auto a_dram_tile = ck_tile::make_tile_scatter_gather( + // a_dram_block_window_tmp.get_bottom_tensor_view(), + // a_dram_block_window_tmp.get_window_lengths(), + // a_dram_block_window_tmp.get_window_origin(), + // a_dist, + // k_offsets); // K DRAM tile window for + + // auto a_copy_dram_window = + // make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), + // make_tuple(number{}, number{}), + // a_dram_block_window_tmp.get_window_origin(), + // PipelinePolicy::template MakeADramTileDistribution()); + + // A LDS tile window for store + auto a_copy_lds_window = make_tile_window( + a_lds_block, make_tuple(number{}, number{}), {0, 0}); + + // A LDS tile for block GEMM + auto a_lds_gemm_window = make_tile_window( + a_lds_block, make_tuple(number{}, number{}), {0, 0}); + + // Block GEMM + auto block_flatmm = BlockFlatmm(); + + // B flat DRAM window for load + auto b_flat_distribution = + PipelinePolicy::template MakeBFlatDramTileDistribution(); + auto b_flat_dram_window = // tile_window_with_static_distribution + make_tile_window( + b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views + make_tuple(number{}, number{}), + b_flat_dram_block_window_tmp.get_window_origin(), + b_flat_distribution); + + // Acc register tile + auto c_block_tile = decltype(block_flatmm(a_lds_gemm_window, b_flat_dram_window)){}; + + // prefetch + // global read 0 + auto a_block_tile = a_dram_block_window.load(); + + { + // move to 1 + move_tile_window(a_dram_block_window, {0, kKPerBlock}); + + // initialize C + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + + // LDS write 0 + if constexpr(std::is_same_v) + { + auto a_shuffle_tmp = make_static_distributed_tensor( + PipelinePolicy::template MakeShuffledARegBlockDistribution()); + shuffle_tile(a_shuffle_tmp, a_block_tile); + const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_shuffle_tmp); + store_tile(a_copy_lds_window, a_block_tile_tmp); + } + else + { + store_tile(a_copy_lds_window, tile_elementwise_in(a_element_func, a_block_tile)); + } + } + + index_t iCounter = num_loop - 1; + while(iCounter > 0) + { + // global read i + 1 + a_dram_block_window.load(a_block_tile); + + block_sync_lds(); + + // GEMM i + block_flatmm(c_block_tile, a_lds_gemm_window, b_flat_dram_window); + + block_sync_lds(); + + // move to i + 2 + move_tile_window(a_dram_block_window, {0, kKPerBlock}); + + // LDS write i + 1 + const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); + store_tile(a_copy_lds_window, a_block_tile_tmp); + + // move to next flat K + move_tile_window(b_flat_dram_window, {0, BlockGemmShape::flatKPerBlock}); + + iCounter--; + } + + // tail + { + block_sync_lds(); + + // GEMM num_loop - 1 + block_flatmm(c_block_tile, a_lds_gemm_window, b_flat_dram_window); + } + + return c_block_tile; + } + + template + CK_TILE_DEVICE auto operator()(ADramBlockWindow& a_dram_block_window_tmp, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()( + a_dram_block_window_tmp, + [](const ADataType& a) { return a; }, + b_flat_dram_block_window_tmp, + num_loop, + p_smem); + } + +}; +} // namespace ck_tile + diff --git a/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp new file mode 100644 index 0000000000..e1924d8d03 --- /dev/null +++ b/include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common.hpp" +#include "ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp" + +namespace ck_tile { +struct MoeGemmPipelineAgBgCrPolicy : public UniversalFlatmmPipelineAgBgCrPolicy +{ + template + CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_C() + { + using S_ = remove_cvref_t; + using WarpGemm = remove_cvref_t; + // using CDataType = typename WarpGemm::CDataType; + + constexpr auto c_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, + sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); + constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); + return c_block_dstr; + } +}; +} // namespace ck_tile + +