From c585cc1429448f1eef49c3bde4e483b6119e7603 Mon Sep 17 00:00:00 2001 From: lalala-sh Date: Wed, 16 Jul 2025 10:40:08 +0800 Subject: [PATCH] support masked mode --- .../19_grouped_flatmm/grouped_flatmm.cpp | 28 +++ .../run_grouped_flatmm_example.inc | 206 ++++++++++++++++++ .../flatmm/kernel/grouped_flatmm_kernel.hpp | 115 ++++++++++ 3 files changed, 349 insertions(+) diff --git a/example/ck_tile/19_grouped_flatmm/grouped_flatmm.cpp b/example/ck_tile/19_grouped_flatmm/grouped_flatmm.cpp index 5c45e37d73..e49faf004f 100644 --- a/example/ck_tile/19_grouped_flatmm/grouped_flatmm.cpp +++ b/example/ck_tile/19_grouped_flatmm/grouped_flatmm.cpp @@ -245,6 +245,34 @@ int run_grouped_flatmm_example(int argc, char* argv[]) throw std::runtime_error("Unsupported data_type!"); } } + else if(mode == "masked") + { + + if(data_type == "fp16") + { + run_masked_grouped_flatmm_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "bf16") + { + run_masked_grouped_flatmm_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "fp8") + { + run_masked_grouped_flatmm_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "bf8") + { + run_masked_grouped_flatmm_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}); + } + else + { + throw std::runtime_error("Unsupported data_type!"); + } + } else { throw std::runtime_error("Unsupported mode!"); diff --git a/example/ck_tile/19_grouped_flatmm/run_grouped_flatmm_example.inc b/example/ck_tile/19_grouped_flatmm/run_grouped_flatmm_example.inc index 2d72b373bd..955769dab0 100644 --- a/example/ck_tile/19_grouped_flatmm/run_grouped_flatmm_example.inc +++ b/example/ck_tile/19_grouped_flatmm/run_grouped_flatmm_example.inc @@ -114,6 +114,35 @@ float invoke_gemm(int n_warmup, int n_repeat, const ck_tile::ContiguousGroupedFl return ave_time; } +template +float invoke_gemm(int n_warmup, int n_repeat, const ck_tile::MaskedGroupedFlatmmHostArgs& args) +{ + float ave_time = + grouped_flatmm( + args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + + std::string op_name{"Grouped Gemm"}; + + std::size_t flop = std::size_t(2) * args.Max_M * args.N * args.K; + std::size_t num_byte = sizeof(ADataType) * args.Max_M * args.K + + sizeof(BDataType) * args.N * args.K + + sizeof(CDataType) * args.Max_M * args.N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_byte / 1.E6 / ave_time; + + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " + << gb_per_sec << " GB/s, " << op_name << std::endl; + + return ave_time; +} + template int run_grouped_flatmm_example_with_layouts(int argc, char* argv[], @@ -527,3 +556,180 @@ int run_contiguous_grouped_flatmm_example_with_layouts( return pass; } + +template +int run_masked_grouped_flatmm_example_with_layouts( + int argc, + char* argv[], + const ALayout a_layout = ALayout{}, + const BLayout b_layout = BLayout{}, + [[maybe_unused]] const CLayout c_layout = CLayout{}) +{ + auto [result, arg_parser] = create_args(argc, argv); + + if(!result) + { + return -1; + }; + + using ADataType = typename GemmBasicTypeConfig::ADataType; + using BDataType = typename GemmBasicTypeConfig::BDataType; + using CDataType = typename GemmBasicTypeConfig::CDataType; + using AccDataType = typename GemmBasicTypeConfig::AccDataType; + + constexpr int BlockM = GemmConfig::M_Tile; + + const int group_count = arg_parser.get_int("group_count"); + const int repeat = arg_parser.get_int("repeat"); + const int warmup = arg_parser.get_int("warmup"); + + std::vector Ms = arg_parser.get_int_vec("Ms"); + std::vector Ns = arg_parser.get_int_vec("Ns"); + std::vector Ks = arg_parser.get_int_vec("Ks"); + + if(!(int(Ms.size()) == group_count)) + { + std::cout << "Please check the input data." << std::endl; + // padding additional Ms if needed + for(int i = 0; i < group_count; i++) + { + Ms.push_back(256 + 64 * i); + } + } + + ck_tile::index_t M = 4096;//Ms[0]; + ck_tile::index_t N = Ns[0]; + ck_tile::index_t K = Ks[0]; + + ck_tile::index_t kbatch = arg_parser.get_int("split_k"); + + ck_tile::index_t stride_A = K; + ck_tile::index_t stride_B = K; + ck_tile::index_t stride_C = N; + + stride_A = ck_tile::get_default_stride(group_count * M, K, stride_A, is_row_major(a_layout)); + stride_B = ck_tile::get_default_stride(K, N * group_count, stride_B, is_row_major(b_layout)); + stride_C = ck_tile::get_default_stride(group_count * M, N, stride_C, is_row_major(c_layout)); + + ck_tile::HostTensor a_m_k_tensor( + ck_tile::host_tensor_descriptor(group_count * M, K, stride_A, is_row_major(a_layout))); + ck_tile::HostTensor b_k_n_tensor(ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(K, N * group_count, stride_B, is_row_major(b_layout)))); + ck_tile::HostTensor c_m_n_tensor(ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(group_count * M, N, stride_C, is_row_major(c_layout)))); + + std::vector m_indices(group_count); + int indices_fill_start = 0; + for(int i = 0; i < group_count; ++i) + { + int group_m = Ms[i]; + int padded_group_m = (group_m + BlockM - 1) / BlockM * BlockM; + for(int j = 0; j < padded_group_m; j++) + { + m_indices[i] = padded_group_m; // -1 for padding + } + } + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k_tensor); + ck_tile::FillUniformDistribution{-.5f, .5f}(b_k_n_tensor); + + constexpr int N_Warp_Tile = GemmConfig::N_Warp_Tile; + assert(N % N_Warp_Tile == 0 && + "N must be divisible by N_Warp_Tile for contiguous grouped gemm"); + ck_tile::HostTensor b_shuffle_host = shuffle_b(b_k_n_tensor); + + std::unique_ptr a_m_k_dev_buf( + std::make_unique(a_m_k_tensor.get_element_space_size_in_bytes())); + std::unique_ptr b_shfl_dev_buf( + std::make_unique(b_shuffle_host.get_element_space_size_in_bytes())); + std::unique_ptr c_m_n_dev_buf( + std::make_unique(c_m_n_tensor.get_element_space_size_in_bytes())); + c_m_n_dev_buf->SetZero(); + + ck_tile::DeviceMem m_indices_dev_buf(group_count * sizeof(ck_tile::index_t)); + m_indices_dev_buf.ToDevice(m_indices.data()); + + a_m_k_dev_buf->ToDevice(a_m_k_tensor.data()); + b_shfl_dev_buf->ToDevice(b_shuffle_host.data()); + + ck_tile::MaskedGroupedFlatmmHostArgs kernal_args{ + static_cast(m_indices_dev_buf.GetDeviceBuffer()), + group_count, + M, + N, + K, + a_m_k_dev_buf->GetDeviceBuffer(), + stride_A, + b_shfl_dev_buf->GetDeviceBuffer(), + stride_B, + c_m_n_dev_buf->GetDeviceBuffer(), + stride_C, + kbatch, + }; + + invoke_gemm( + warmup, repeat, kernal_args); + c_m_n_dev_buf->FromDevice(c_m_n_tensor.data()); + + bool pass{true}; + if(arg_parser.get_int("v") == 1) + { + throw std::runtime_error( + "Not support v=1 host verification in contiguous grouped gemm, use " + "v=2 device verification instead"); + } + else if(arg_parser.get_int("v") == 2) + { + BDataType* d_B; + CDataType* d_C; + ck_tile::hip_check_error(hipMalloc(&d_B, N * K * sizeof(BDataType))); + ck_tile::hip_check_error(hipMalloc(&d_C, M * N * sizeof(CDataType))); + ck_tile::hip_check_error(hipMemset(d_C, 0, M * N * sizeof(CDataType))); + + ck_tile::HostTensor c_gpu_ref_host( + ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); + + ck_tile::index_t acc_m = 0; + for(int i = 0; i < group_count; ++i) + { + ck_tile::hip_check_error(hipMemcpy(d_B, + b_k_n_tensor.data() + i * N * K, + N * K * sizeof(BDataType), + hipMemcpyHostToDevice)); + ck_tile::reference_gemm_gpu( + static_cast(a_m_k_dev_buf->GetDeviceBuffer()) + i * M * K, + d_B, + d_C + i * M * N, + m_indices[i], + N, + K, + stride_A, + stride_B, + stride_C); + ck_tile::hip_check_error(hipMemcpy( + c_gpu_ref_host.data(), d_C, m_indices[i] * N * sizeof(CDataType), hipMemcpyDeviceToHost)); + } + + + ck_tile::hip_check_error(hipFree(d_B)); + ck_tile::hip_check_error(hipFree(d_C)); + + float rtol = 1e-3; + float atol = 1e-3; + + pass = ck_tile::check_err( + c_m_n_tensor, c_gpu_ref_host, "Error: Incorrect results!", rtol, atol); + + std::cout << "Relative error threshold: " << rtol << " Absolute error threshold: " << atol + << std::endl; + std::cout << "The GPU veification result is: " << (pass ? "correct" : "fail") << std::endl; + } + + return pass; +} diff --git a/include/ck_tile/ops/flatmm/kernel/grouped_flatmm_kernel.hpp b/include/ck_tile/ops/flatmm/kernel/grouped_flatmm_kernel.hpp index 683ebb91ec..604ae08256 100644 --- a/include/ck_tile/ops/flatmm/kernel/grouped_flatmm_kernel.hpp +++ b/include/ck_tile/ops/flatmm/kernel/grouped_flatmm_kernel.hpp @@ -94,6 +94,50 @@ struct ContiguousGroupedFlatmmHostArgs index_t k_batch; }; +struct MaskedGroupedFlatmmHostArgs +{ + CK_TILE_HOST MaskedGroupedFlatmmHostArgs() = default; + CK_TILE_HOST MaskedGroupedFlatmmHostArgs(index_t* M_indices_, + index_t group_count_, + index_t Max_M_, + index_t N_, + index_t K_, + const void* a_ptr_, + index_t stride_A_, + const void* b_shuffle_ptr_, + index_t stride_B_, + void* c_ptr_, + index_t stride_C_, + index_t k_batch_) + : M_indices(M_indices_), + group_count(group_count_), + Max_M(Max_M_), + N(N_), + K(K_), + a_ptr(a_ptr_), + stride_A(stride_A_), + b_shuffle_ptr(b_shuffle_ptr_), + stride_B(stride_B_), + c_ptr(c_ptr_), + stride_C(stride_C_), + k_batch(k_batch_) + { + } + + index_t* M_indices; + index_t group_count; + index_t Max_M; + index_t N; + index_t K; + const void* a_ptr; + index_t stride_A; + const void* b_shuffle_ptr; + index_t stride_B; + void* c_ptr; + index_t stride_C; + index_t k_batch; +}; + template struct GroupedFlatmmKernel : FlatmmKernel { @@ -174,6 +218,35 @@ struct GroupedFlatmmKernel : FlatmmKernel( + kentry2), + block_size, + dync_smem_size); + + const int persistent_block_size = prop.multiProcessorCount * maxActiveBlocksPerCU; + // const int total_work_tile_cnt = TilePartitioner::GridSize(kernelArgs.M, kernelArgs.N); + + std::cout << "maxActiveBlocksPerCU: " << maxActiveBlocksPerCU + << ", persistent_block_size: " << persistent_block_size << std::endl; + + assert(kernelArgs.k_batch == 1); + return dim3(persistent_block_size, 1, kernelArgs.k_batch); + } + CK_TILE_HOST static constexpr auto MakeKernelArgs(const GroupedFlatmmHostArgs& hostArgs) { return hostArgs; @@ -183,6 +256,11 @@ struct GroupedFlatmmKernel : FlatmmKernel(kargs.a_ptr) + group_idx * kargs.Max_M * kargs.K, + static_cast(kargs.b_shuffle_ptr) + group_idx * kargs.N * kargs.K, + static_cast(kargs.c_ptr) + group_idx * kargs.Max_M * kargs.N, + M, + kargs.N, + kargs.K, + kargs.stride_A, + kargs.stride_B, + kargs.stride_C, + kargs.k_batch, + }; + // call the underlying flatmm kernel + underlying_kernel(impl_kargs, block_linear_idx); + block_linear_idx += total_block_cnt; + } + block_linear_idx -= group_block_cnt; + } + } }; } // namespace ck_tile