From 30ab1d6a7108b6f9b4463f8e8183e223428222c0 Mon Sep 17 00:00:00 2001 From: Mateusz Ozga <110818320+mozga-amd@users.noreply.github.com> Date: Fri, 19 Sep 2025 01:14:11 +0200 Subject: [PATCH 01/12] [CK_TILE] Multiple-ABD GEMM example (#2788) * Multi ABD - initial commit * Clang-foramt fix * block gemm, unify the name of CDataType * Apply chnages to mem-pipeline * Rollback prefix for DType and Layout * Gemm Kernel Basic, rename * WMMA config * Grouped GEMM * Clang-format * Dropout, name * Review v2 * Move element_wise fn to unnary, remov old ones fn * clang-format * Fix issue review * WP operator adjust to universal gemm * v2 prepare * Remove unused comment * Remove vectorsize * Rollback * Adjust pipeline for abd * Shuffle argument * CI-fail fix quant * Fix ag_br pipeline * Failing tests * Typo * Single argument support --- CHANGELOG.md | 1 + .../ck_tile/22_gemm_multi_abd/CMakeLists.txt | 1 + example/ck_tile/22_gemm_multi_abd/README.md | 35 ++ .../22_gemm_multi_abd/gemm_multi_abd_fp16.cpp | 184 +++++++ .../22_gemm_multi_abd/gemm_multi_abd_fp16.hpp | 186 +++++++ .../run_gemm_multi_abd_fp16_example.inc | 311 +++++++++++ example/ck_tile/22_gemm_multi_abd/utils.hpp | 38 ++ example/ck_tile/CMakeLists.txt | 1 + include/ck_tile/core/tensor/load_tile.hpp | 23 + include/ck_tile/core/tensor/tile_window.hpp | 143 +++++ .../ck_tile/host/reference/reference_gemm.hpp | 75 +++ .../unary_element_wise_operation.hpp | 17 + .../ops/epilogue/cshuffle_epilogue.hpp | 27 +- .../ops/epilogue/default_2d_epilogue.hpp | 32 +- include/ck_tile/ops/gemm.hpp | 1 + .../ops/gemm/kernel/batched_gemm_kernel.hpp | 4 +- .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 8 +- .../ops/gemm/kernel/gemm_multi_abd_kernel.hpp | 193 +++++++ .../ops/gemm/kernel/gemm_multi_d_kernel.hpp | 8 +- .../ops/gemm/kernel/grouped_gemm_kernel.hpp | 12 +- .../ops/gemm/kernel/universal_gemm_kernel.hpp | 42 +- .../pipeline/gemm_pipeline_ag_bg_cr_base.hpp | 129 ++++- .../gemm_pipeline_ag_bg_cr_comp_v3.hpp | 253 ++++++--- .../gemm_pipeline_ag_bg_cr_comp_v4.hpp | 311 +++++++---- .../gemm_pipeline_ag_bg_cr_comp_v5.hpp | 146 +++-- .../pipeline/gemm_pipeline_ag_bg_cr_mem.hpp | 320 +++++++---- .../gemm_pipeline_agmem_bgmem_creg_v1.hpp | 152 ++++-- .../gemm_pipeline_agmem_bgmem_creg_v2.hpp | 164 ++++-- .../gemm/pipeline/gemm_pipeline_problem.hpp | 137 +++-- ...emm_universal_pipeline_ag_bg_cr_policy.hpp | 28 +- .../ops/gemm/pipeline/tile_gemm_traits.hpp | 28 +- .../wp_pipeline_agmem_bgmem_creg_v1.hpp | 60 ++- .../wp_pipeline_agmem_bgmem_creg_v2.hpp | 65 ++- .../pipeline/tile_gemm_quant_traits.hpp | 4 + test/ck_tile/CMakeLists.txt | 1 + test/ck_tile/gemm_multi_abd/CMakeLists.txt | 12 + .../test_gemm_multi_abd_cshuffle.cpp | 40 ++ .../test_gemm_multi_abd_default2d.cpp | 41 ++ .../test_gemm_multi_abd_ut_cases_cshuffle.inc | 211 ++++++++ ...test_gemm_multi_abd_ut_cases_default2d.inc | 211 ++++++++ .../test_gemm_multi_abd_util.hpp | 500 ++++++++++++++++++ 41 files changed, 3603 insertions(+), 552 deletions(-) create mode 100644 example/ck_tile/22_gemm_multi_abd/CMakeLists.txt create mode 100644 example/ck_tile/22_gemm_multi_abd/README.md create mode 100644 example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp create mode 100644 example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp create mode 100644 example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc create mode 100644 example/ck_tile/22_gemm_multi_abd/utils.hpp create mode 100644 include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp create mode 100644 test/ck_tile/gemm_multi_abd/CMakeLists.txt create mode 100644 test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_cshuffle.cpp create mode 100644 test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_default2d.cpp create mode 100644 test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_cshuffle.inc create mode 100644 test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_default2d.inc create mode 100644 test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_util.hpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 38669385f3..dafe1b5c87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj * Added support for GKCYX layout for grouped convolution backward data (NGCHW/GKCYX/NGKHW). * Added support for Stream-K version of mixed fp8/bf16 GEMM * Added support for Multiple D GEMM +* Added support for Multiple ABD GEMM * Added GEMM pipeline for microscaling (MX) FP8/FP6/FP4 data types * Added support for FP16 2:4 structured sparsity to universal GEMM. * Added support for Split K for grouped convolution backward data. diff --git a/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt b/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt new file mode 100644 index 0000000000..f382e0cf45 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt @@ -0,0 +1 @@ +add_executable(tile_example_gemm_multi_abd_fp16 EXCLUDE_FROM_ALL gemm_multi_abd_fp16.cpp) diff --git a/example/ck_tile/22_gemm_multi_abd/README.md b/example/ck_tile/22_gemm_multi_abd/README.md new file mode 100644 index 0000000000..c272df3fb5 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/README.md @@ -0,0 +1,35 @@ +#Multiple ABD GEMM + +This folder contains example for Multiple ABD GEMM using ck_tile tile-programming implementation. + +## build +``` +#in the root of ck_tile +mkdir build && cd build +#you can replace < arch> with the appropriate architecture(for example gfx90a or gfx942) or \ + leave it blank +sh ../script/cmake-ck-dev.sh ../ +#The basic pipeline method on the gemm calculation +make tile_example_gemm_multi_abd_fp16 -j +``` +This will result in an executable `build/bin/tile_example_gemm_multi_abd_fp16` + +## example +``` +args: + -m M dimensions - (Default: 3840) + -n N dimensions - (Default: 4096) + -k K dimensions - (Default: 4096) +-as_layout Tensor A layout (default:R) +-bs_layout Tensor B layout (default:C) +-ds_layout Tensor D layout (default:R) +-e_layout Tensor E layout (default:R) +-stride_as Tensor A strides - (Default: 0) +-stride_bs Tensor B strides - (Default: 0) +-stride_e Tensor C strides - (Default: 0) +-stride_ds Tensor D strides - (Default: 0) +-validate 0. No validation, 1. Validation on GPU. (Default: 1) + -warmup Number of iterations before benchmark the kernel. (Default: 10) + -repeat Number of iterations to benchmark the kernel. (Default: 100) + -kbatch kbatch for SplitK. (Default: 1) +``` \ No newline at end of file diff --git a/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp new file mode 100644 index 0000000000..6d955c3a09 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/host.hpp" +#include "gemm_multi_abd_fp16.hpp" +#include "utils.hpp" + +template +auto gemm_multi_abd(const gemm_multi_abd_kargs& args, const ck_tile::stream_config& s) -> float +{ + constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; + constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; + constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; + + constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; + constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; + constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; + + constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; + constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; + constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; + + constexpr bool DoubleSmemBuffer = GemmConfig::DoubleSmemBuffer; + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + + constexpr bool TransposeC = false; + + constexpr int kBlockPerCu = 1; + constexpr ck_tile::index_t TileParitionerGroupNum = 8; + constexpr ck_tile::index_t TileParitionerM01 = 4; + + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; + + using Traits = ck_tile::TileGemmTraits; + + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmPipelineProblem = + ck_tile::GemmPipelineProblem; + + using BaseGemmPipeline = typename PipelineTypeTraits< + GemmConfig::Pipeline>::template UniversalGemmPipeline; + + const ck_tile::index_t k_grain = args.k_batch * K_Tile; + const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + float ave_time{0}; + + const auto Run = + [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr auto scheduler = GemmConfig::Scheduler; + constexpr auto memory_operation = memory_operation_.value; + + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + + using GemmPipeline = typename PipelineTypeTraits< + GemmConfig::Pipeline>::template GemmPipeline; + + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem>; + + using Kernel = ck_tile::GemmKernelMultiABD; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args:" << " grid: {" << grids.x << ", " + << grids.y << ", " << grids.z << "}" << ", blocks: {" << blocks.x << ", " + << blocks.y << ", " << blocks.z << "}" << std::endl; + } + + ave_time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + return ave_time; + }; + + const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) { + if(args.k_batch == 1) + { + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + else + { + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + }; + + BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num); + + return ave_time; +} + +#include "run_gemm_multi_abd_fp16_example.inc" + +int main(int argc, char* argv[]) +{ +#if CK_TILE_USE_WMMA + return !run_multiple_abd_gemm_example(argc, argv); +#else + return !run_multiple_abd_gemm_example(argc, argv); +#endif +} diff --git a/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp new file mode 100644 index 0000000000..35bc232eca --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" + +#define CK_TILE_PIPELINE_COMPUTE_V3 1 +#define CK_TILE_PIPELINE_MEMORY 2 +#define CK_TILE_PIPELINE_COMPUTE_V4 3 + +#ifndef CK_TILE_PIPELINE_DEFAULT +#define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 +#endif + +using A0DataType = ck_tile::half_t; +using A1DataType = ck_tile::half_t; + +using B0DataType = ck_tile::half_t; +using B1DataType = ck_tile::half_t; + +using D0DataType = ck_tile::half_t; +using D1DataType = ck_tile::half_t; + +using EDataType = ck_tile::half_t; + +using AsDataType = ck_tile::tuple; +using BsDataType = ck_tile::tuple; +using DsDataType = ck_tile::tuple; + +using AccDataType = float; + +struct GemmConfigMemory +{ + // Memory friendly for Interwave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 32; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 4; + static constexpr ck_tile::index_t N_Warp = 1; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 8; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; +}; + +struct GemmConfigV3 +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +struct GemmConfigV4 +{ + // Compute friendly for Intrawave scheduler + // Using the ping pong reader in the lds level + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 32; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = true; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +struct GemmConfigV3_Wmma +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +template +struct PipelineTypeTraits; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem; +}; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; +}; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV4; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV4; +}; + +auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "m dimension") + .insert("n", "4096", "n dimension") + .insert("k", "4096", "k dimension") + .insert("as_layout", "R", "As tensor data layout - Row by default") + .insert("bs_layout", "C", "Bs tensor data layout - Col by default") + .insert("ds_layout", "R", "Ds tensor data layout - Row by default") + .insert("e_layout", "R", "E tensor data layout - Row by default") + .insert("stride_as", "0", "Tensor A stride") + .insert("stride_bs", "0", "Tensor B stride") + .insert("stride_ds", "0", "Tensor Ds stride") + .insert("stride_e", "0", "Tensor E stride") + .insert("v", "1", "0. No validation, 1. Validation on GPU") + .insert("warmup", "50", "number of iterations before benchmark the kernel") + .insert("repeat", "100", "number of iterations to benchmark the kernel") + .insert("kbatch", "1", "kbatch for SplitK"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} +using gemm_multi_abd_kargs = + ck_tile::GemmMultiABDHostArgs; + +template +float gemm_multi_abd(const gemm_multi_abd_kargs& kargs, const ck_tile::stream_config& s); diff --git a/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc b/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc new file mode 100644 index 0000000000..881961c9db --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc @@ -0,0 +1,311 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once +#include + +template +float invoke_gemm_multi_abd(const std::array& as_m_k_dev_buf, + const std::array& bs_k_n_dev_buf, + const std::array& ds_m_n_dev_buf, + void* e_m_n_dev_buf, + ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + const std::array& StrideAs, + const std::array& StrideBs, + const std::array& StrideDs, + ck_tile::index_t StrideE, + int n_warmup, + int n_repeat, + int k_batch) +{ + gemm_multi_abd_kargs gemm_descs({as_m_k_dev_buf, + bs_k_n_dev_buf, + ds_m_n_dev_buf, + e_m_n_dev_buf, + k_batch, + M, + N, + K, + StrideAs, + StrideBs, + StrideDs, + StrideE}); + + float ave_time = gemm_multi_abd( + gemm_descs, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + + std::string op_name{"Gemm Multiple-ABD"}; + + std::size_t flop = 0, num_btype = 0; + + flop += std::size_t(2) * M * N * K; + + num_btype += + sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(EDataType) * M * N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Run Gemm Multiple-ABD kernel with:\n"; + std::cout << "M =" << M << " N =" << N << " K =" << K << "\n"; + std::cout << "StrideA = " << StrideAs[0] << " StrideB = " << StrideBs[0] + << " StrideE = " << StrideE << "\n"; + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << "\n"; + + return ave_time; +} + +template +int run_gemm_multi_abd_example_with_layouts(int argc, + char* argv[], + const A0Layout a0_layout = A0Layout{}, + const A1Layout a1_layout = A1Layout{}, + const B0Layout b0_layout = B0Layout{}, + const B1Layout b1_layout = B1Layout{}, + const D0Layout d0_layout = D0Layout{}, + const D1Layout d1_layout = D1Layout{}, + const ELayout e_layout = ELayout{}) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + using AElementWiseFn = ck_tile::element_wise::AddScale; + using BElementWiseFn = ck_tile::element_wise::AddScale; + using CDEElementWiseFn = ck_tile::element_wise::MultiDMultiply; + using AsLayout = ck_tile::tuple; + using BsLayout = ck_tile::tuple; + using DsLayout = ck_tile::tuple; + + ck_tile::index_t M = arg_parser.get_int("m"); + ck_tile::index_t N = arg_parser.get_int("n"); + ck_tile::index_t K = arg_parser.get_int("k"); + + ck_tile::index_t StrideA = arg_parser.get_int("stride_as"); + ck_tile::index_t StrideB = arg_parser.get_int("stride_bs"); + ck_tile::index_t StrideD = arg_parser.get_int("stride_ds"); + ck_tile::index_t StrideE = arg_parser.get_int("stride_e"); + + ck_tile::index_t StrideA0 = StrideA; + ck_tile::index_t StrideA1 = StrideA; + + ck_tile::index_t StrideB0 = StrideB; + ck_tile::index_t StrideB1 = StrideB; + + ck_tile::index_t StrideD0 = StrideD; + ck_tile::index_t StrideD1 = StrideD; + + const int n_warmup = arg_parser.get_int("warmup"); + const int n_repeat = arg_parser.get_int("repeat"); + const int k_batch = arg_parser.get_int("kbatch"); + + StrideA0 = get_default_stride(M, N, StrideA0, is_row_major(a1_layout)); + StrideA1 = get_default_stride(M, N, StrideA1, is_row_major(a1_layout)); + + StrideB0 = get_default_stride(K, N, StrideB0, is_row_major(b0_layout)); + StrideB1 = get_default_stride(K, N, StrideB1, is_row_major(b1_layout)); + + StrideD0 = get_default_stride(M, N, StrideD0, is_row_major(d0_layout)); + StrideD1 = get_default_stride(M, N, StrideD1, is_row_major(d1_layout)); + + StrideE = get_default_stride(M, N, StrideE, is_row_major(e_layout)); + + ck_tile::HostTensor a0_m_k_tesnor( + host_tensor_descriptor(M, K, StrideA0, is_row_major(a0_layout))); + ck_tile::HostTensor a1_m_k_tesnor( + host_tensor_descriptor(M, K, StrideA1, is_row_major(a1_layout))); + + ck_tile::HostTensor b0_k_n_tensors( + host_tensor_descriptor(K, N, StrideB0, is_row_major(b0_layout))); + ck_tile::HostTensor b1_k_n_tensors( + host_tensor_descriptor(K, N, StrideB1, is_row_major(b1_layout))); + + ck_tile::HostTensor d0_m_n_tensors( + host_tensor_descriptor(M, N, StrideD0, is_row_major(d0_layout))); + ck_tile::HostTensor d1_m_n_tensors( + host_tensor_descriptor(M, N, StrideD1, is_row_major(d1_layout))); + + ck_tile::HostTensor e_m_n_device_result( + host_tensor_descriptor(M, N, StrideE, is_row_major(e_layout))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a0_m_k_tesnor); + ck_tile::FillUniformDistribution{-1.f, 1.f}(a1_m_k_tesnor); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(b0_k_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b1_k_n_tensors); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(d0_m_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(d1_m_n_tensors); + + ck_tile::DeviceMem a0_m_k_dev_buf(a0_m_k_tesnor.get_element_space_size_in_bytes()); + ck_tile::DeviceMem a1_m_k_dev_buf(a1_m_k_tesnor.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem b0_k_n_dev_buf(b0_k_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b1_k_n_dev_buf(b1_k_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem d0_m_n_dev_buf(d0_m_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d1_m_n_dev_buf(d1_m_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem e_m_n_dev_buf(e_m_n_device_result.get_element_space_size_in_bytes()); + + a0_m_k_dev_buf.ToDevice(a0_m_k_tesnor.mData.data()); + a1_m_k_dev_buf.ToDevice(a1_m_k_tesnor.mData.data()); + + b0_k_n_dev_buf.ToDevice(b0_k_n_tensors.mData.data()); + b1_k_n_dev_buf.ToDevice(b1_k_n_tensors.mData.data()); + + d0_m_n_dev_buf.ToDevice(d0_m_n_tensors.mData.data()); + d1_m_n_dev_buf.ToDevice(d1_m_n_tensors.mData.data()); + + e_m_n_dev_buf.SetZero(); + e_m_n_device_result.SetZero(); + + std::array as_ptr_buf = {a0_m_k_dev_buf.GetDeviceBuffer(), + a1_m_k_dev_buf.GetDeviceBuffer()}; + + std::array bs_ptr_buf = {b0_k_n_dev_buf.GetDeviceBuffer(), + b1_k_n_dev_buf.GetDeviceBuffer()}; + + std::array ds_ptr_buf = {d0_m_n_dev_buf.GetDeviceBuffer(), + d1_m_n_dev_buf.GetDeviceBuffer()}; + + std::array strideAs = {StrideA0, StrideA1}; + std::array strideBs = {StrideB0, StrideB1}; + std::array strideDs = {StrideD0, StrideD1}; + + invoke_gemm_multi_abd(as_ptr_buf, + bs_ptr_buf, + ds_ptr_buf, + e_m_n_dev_buf.GetDeviceBuffer(), + M, + N, + K, + strideAs, + strideBs, + strideDs, + StrideE, + n_warmup, + n_repeat, + k_batch); + + e_m_n_dev_buf.FromDevice(e_m_n_device_result.data()); + + ck_tile::HostTensor a_m_k_host_ref_element_result( + host_tensor_descriptor(M, K, StrideA0, is_row_major(a0_layout))); + ck_tile::HostTensor b_k_n_host_ref_element_result( + host_tensor_descriptor(K, N, StrideB0, is_row_major(b0_layout))); + ck_tile::HostTensor e_m_n_host_ref( + host_tensor_descriptor(M, N, StrideE, is_row_major(e_layout))); + a_m_k_host_ref_element_result.SetZero(); + b_k_n_host_ref_element_result.SetZero(); + e_m_n_host_ref.SetZero(); + + ck_tile::reference_gemm_multiple_abd({a0_m_k_tesnor, a1_m_k_tesnor}, + {b0_k_n_tensors, b1_k_n_tensors}, + {d0_m_n_tensors, d1_m_n_tensors}, + a_m_k_host_ref_element_result, + b_k_n_host_ref_element_result, + e_m_n_host_ref); + + bool pass{true}; + if(arg_parser.get_int("v")) + { + const float max_accumulated_value = + *std::max_element(e_m_n_host_ref.mData.begin(), e_m_n_host_ref.mData.end()); + + const auto rtol_atol = calculate_rtol_atol(K, 1, max_accumulated_value); + + pass &= ck_tile::check_err(e_m_n_device_result, + e_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) + << std::endl; + std::cout << "Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + std::cout << "The CPU veification result is: " << (pass ? "correct" : "fail") << std::endl; + } + return pass; +} + +template +int run_multiple_abd_gemm_example(int argc, char* argv[]) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + + const std::string as_layout = arg_parser.get_str("as_layout"); + const std::string bs_layout = arg_parser.get_str("bs_layout"); + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + + if(as_layout == "R" && bs_layout == "C") + { + return run_gemm_multi_abd_example_with_layouts( + argc, argv, Row{}, Row{}, Col{}, Col{}, Row{}, Row{}, Row{}); + } + else + { + throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!"); + } +} diff --git a/example/ck_tile/22_gemm_multi_abd/utils.hpp b/example/ck_tile/22_gemm_multi_abd/utils.hpp new file mode 100644 index 0000000000..38bf8623d4 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/utils.hpp @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +template +static constexpr inline auto is_row_major(Layout layout_) +{ + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeTypeAB = + std::conditional_t; + + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} diff --git a/example/ck_tile/CMakeLists.txt b/example/ck_tile/CMakeLists.txt index 8fce70ba04..75d32a5eb0 100644 --- a/example/ck_tile/CMakeLists.txt +++ b/example/ck_tile/CMakeLists.txt @@ -21,6 +21,7 @@ add_subdirectory(18_flatmm) add_subdirectory(19_gemm_multi_d) add_subdirectory(20_grouped_convolution) add_subdirectory(21_elementwise) +add_subdirectory(22_gemm_multi_abd) add_subdirectory(35_batched_transpose) add_subdirectory(38_block_scale_gemm) add_subdirectory(39_copy) diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index 8b7541bf23..c7c4702e22 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -26,6 +26,29 @@ CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window, return tile_window.load(number{}, bool_constant{}); } +/** + * @brief Load tile with elementwise function + * + * @note This function is a modification of the existing load function. + * It has been extended with two additional parameters: it takes a tuple as input + * and an elementwise function. For each A = A0, A1… AN, the elementwise function + * is additionally applied during a single read. + */ +template +CK_TILE_DEVICE auto load_tile_with_elementwise(const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) +{ + // TODO: Tile windows should works with unknow number of params + // Load element_wise API works only when the input typle is a tuple-tyupe + return tile_window[number<0>{}].load( + tile_window, elementwise, number{}, bool_constant{}); +} + template + CK_TILE_DEVICE auto load(const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) const + { + constexpr auto tile_dstr = typename Base::TileDstr{}; + auto dst_tensor = make_static_distributed_tensor(tile_dstr); + load(dst_tensor, + tile_window, + elementwise, + number{}, + bool_constant{}); + return dst_tensor; + } + + template + CK_TILE_DEVICE auto load(DistributedTensor& dst_tensor, + const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) const + { + + using Traits = typename Base::Traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = typename Base::TileDstr{}; + constexpr auto sizeOfTuple = TileWindow_::size(); + // 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 = + tile_window[number<0>{}].pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = + tile_window[number<0>{}].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); + + // read from bottom tensor + const auto idx_vec_value = generate_tuple( + [&](auto jj) { + return tile_window[number{}] + .get_bottom_tensor_view() + .template get_vectorized_elements( + bottom_tensor_thread_coord, + 0, + bool_constant{}); + }, + number{}); + + // 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; + + ck_tile::apply( + [&](auto&&... t) { + elementwise(dst_tensor.get_thread_buffer().template at(), + t.template get_as< + typename Base::DataType>()[j / Traits::PackedSize]...); + }, + idx_vec_value); + }); + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + idx_diff_ys); + + Base::move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + } + }); + }); + } + template @@ -857,6 +967,39 @@ CK_TILE_DEVICE void move_tile_window( window.move(step); } +template +CK_TILE_DEVICE void move_tile_window( + tuple>& window, + const typename tile_window_with_static_distribution::BottomTensorIndex& step) +{ + using T = tuple>; + + static constexpr auto N = T::size(); + static_for<0, N, 1>{}([&](auto Is) { window[number{}].move(step); }); +} + +template ::value>* = nullptr> +CK_TILE_DEVICE void move_tile_window(TileWindowWithStaticDistributionType& window, StepType& step) +{ + static constexpr auto N = TileWindowWithStaticDistributionType::size(); + static_for<0, N, 1>{}([&](auto Is) { window[number{}].move(step); }); +} + /** * @brief This class provides description of tile windowed view on the device memory. * diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index caa00e5994..d9379b4420 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -261,6 +261,81 @@ CK_TILE_HOST void reference_gemm(const HostTensor& a_m_k, make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); } +template >, + typename BDataType = remove_cvref_t>, + typename DDataType = remove_cvref_t>> +CK_TILE_HOST void +reference_gemm_multiple_abd(const std::array, AsDataType::size()>& as_m_k, + const std::array, BsDataType::size()>& bs_k_n, + const std::array, DsDataType::size()>& ds_m_n, + HostTensor& a_m_k, + HostTensor& b_k_n, + HostTensor& c_m_n, + const AElementOp& a_element_op = {}, + const BElementOp& b_element_op = {}, + const CDElementOp& acc_element_op = {}) +{ + const std::size_t M = a_m_k.get_length(0); + const std::size_t N = b_k_n.get_length(1); + const std::size_t K = a_m_k.get_length(1); + + auto as_m_k_tuple = + generate_tie([&](auto idx) -> auto& { return as_m_k[idx]; }, number{}); + + auto bs_k_n_tuple = + generate_tie([&](auto idx) -> auto& { return bs_k_n[idx]; }, number{}); + + auto ds_m_n_tuple = + generate_tie([&](auto idx) -> auto& { return ds_m_n[idx]; }, number{}); + + // Apply elementwise function to A + auto a_elementwise_fn = [&](auto i, auto j) { + ck_tile::apply([&](auto&&... t) { a_element_op(a_m_k(i, j), t(i, j)...); }, as_m_k_tuple); + }; + + make_ParallelTensorFunctor(a_elementwise_fn, M, K)(std::thread::hardware_concurrency()); + + // Apply elementwise function to B + auto b_elementwise_fn = [&](auto i, auto j) { + ck_tile::apply([&](auto&&... t) { b_element_op(b_k_n(i, j), t(i, j)...); }, bs_k_n_tuple); + }; + + make_ParallelTensorFunctor(b_elementwise_fn, K, N)(std::thread::hardware_concurrency()); + + auto f_mk_kn_mn = [&](auto m, auto n) { + AccDataType v_acc = 0; + for(std::size_t k = 0; k < K; ++k) + { + ADataType v_a = a_m_k(m, k); + BDataType v_b = b_k_n(k, n); + v_acc += + ck_tile::type_convert(v_a) * ck_tile::type_convert(v_b); + } + + CDataType v_c = 0; + + ck_tile::apply( + [&](auto&&... t) { + acc_element_op(v_c, + ck_tile::type_convert(v_acc), + ck_tile::type_convert(t(m, n))...); + }, + ds_m_n_tuple); + + c_m_n(m, n) = ck_tile::type_convert(v_c); + }; + + make_ParallelTensorFunctor(f_mk_kn_mn, M, N)(std::thread::hardware_concurrency()); +} + template + CK_TILE_HOST_DEVICE constexpr void operator()(E& a, const As&... as) const + { + // Start with the base value c + float result = ck_tile::type_convert(0.0f); + + // Add by each D parameter using fold expression + ((result += ck_tile::type_convert(as)), ...); + + a = ck_tile::type_convert(scale * result); + } + + float scale = 1.0; +}; + struct MultiDMultiply { template diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 628af0e0b3..ebd97c1c66 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -28,8 +28,8 @@ struct GetDataType using type = typename T::DataType; // Use T::ScaleN::DataType }; -template struct CShuffleEpilogueProblem { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; using DsDataType = remove_cvref_t; @@ -83,12 +83,27 @@ template struct CShuffleEpilogue { using Problem = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; using DsDataType = remove_cvref_t; using DsLayout = remove_cvref_t; + + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using ATypeToUse = std::conditional_t, BDataType, ADataType>; // Used for weight-only quantization kernel, B would be dequantized to the same data type as A diff --git a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp index 54becd3c0f..2843966cd7 100644 --- a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp @@ -28,8 +28,8 @@ struct Default2DEpilogueProblem static constexpr index_t NumDTensor = 0; }; -template { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CLayout = remove_cvref_t; using DsDataType = remove_cvref_t; using CDElementwise = remove_cvref_t; @@ -157,14 +157,28 @@ struct Default2DEpilogue template struct DefaultGemm2DEpilogue : public Default2DEpilogue { - using Problem = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using AccDataType = remove_cvref_t; - using ODataType = remove_cvref_t; + using Problem = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using AccDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; // Used for weight-only quantization kernel, B would be dequantized to the same data type as A using BTypeToUse = std::conditional_t, ADataType, BDataType>; + using DsDataType = remove_cvref_t; using DsLayout = remove_cvref_t; using CDElementwise = remove_cvref_t; diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index de13e305e0..6e07dbc00e 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -31,6 +31,7 @@ #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_multi_d_kernel.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp" diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index fcfbf9635f..588d903b25 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -90,10 +90,10 @@ struct BatchedGemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); struct BatchedGemmKernelArgs : ck_tile::UniversalGemmKernelArgs<> { diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index e37b4f36d4..d632b1596c 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -89,7 +89,7 @@ struct GemmKernel /// @brief Specify the layout configurations for A, B, E and D using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; /// @brief Specify the data type configurations for A, B, E and D using ADataType = remove_cvref_t; @@ -106,10 +106,10 @@ struct GemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. - static_assert(!is_detected::value && + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); static constexpr index_t NumATensor = 1; static constexpr index_t NumBTensor = 1; diff --git a/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp new file mode 100644 index 0000000000..3b050e03ed --- /dev/null +++ b/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common.hpp" +#include "ck_tile/host/concat.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/stream_utils.hpp" +#include "ck_tile/core/utility/env.hpp" +#include "ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp" +#include "ck_tile/core/utility/type_traits.hpp" + +namespace ck_tile { + +/// @brief The MultiABD GEMM kernel host arguments. +/// +/// @par Overview +/// This structure is passed to @ref GemmKernelMultiABD "GemmKernelMultiABD" when creating +/// kernel arguments object. It contain all necessary information required to build proper +/// kernel argument and launch kernel on GPU. This structure defines the GEMM problem +/// configuration by stating all required information like M,N,K sizes and respective strides. +/// NumATensor describes the number of A tensors. The minimum number of tensors is 1(required). +/// NumBTensor describes the number of B tensors. The minimum number of tensors is 1(required). +/// NumDTensor describes the number of D tensors. The minimum number of tensors is 0(not +/// required). +template +struct GemmMultiABDHostArgs +{ + CK_TILE_HOST GemmMultiABDHostArgs(const std::array& as_ptr_, + const std::array& bs_ptr_, + const std::array& ds_ptr_, + void* e_ptr_, + index_t k_batch_, + index_t M_, + index_t N_, + index_t K_, + const std::array& stride_As_, + const std::array& stride_Bs_, + const std::array& stride_Ds_, + index_t stride_E_) + : as_ptr(as_ptr_), + bs_ptr(bs_ptr_), + ds_ptr(ds_ptr_), + e_ptr(e_ptr_), + M(M_), + N(N_), + K(K_), + stride_As(stride_As_), + stride_Bs(stride_Bs_), + stride_Ds(stride_Ds_), + stride_E(stride_E_), + k_batch(k_batch_) + { + } + + const std::array as_ptr; + const std::array bs_ptr; + const std::array ds_ptr; + union + { + void* e_ptr; + void* c_ptr; + }; + index_t M; + index_t N; + index_t K; + const std::array stride_As; + const std::array stride_Bs; + const std::array stride_Ds; + union + { + index_t stride_E; + index_t stride_C; + }; + + index_t k_batch; +}; + +template +struct GemmKernelMultiABD +{ + /// @brief Inject the UniversalGemmKernel base class to support execution of all necessary + /// functions. + using UniversalGemmKernel = + UniversalGemmKernel; + static constexpr index_t kBlockSize = UniversalGemmKernel::kBlockSize; + + using TilePartitioner = remove_cvref_t; + using GemmPipeline = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + + /// @brief Specify the layout configurations for A, B, E and D + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + using DsLayout = remove_cvref_t; + + /// @brief Specify the data type configurations for A, B, E and D + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using EDataType = remove_cvref_t; + using DsDataType = remove_cvref_t; + + /// @brief ALayout and ADataType are expected to be a tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value, + "ALayout and ADataType must be a tuple."); + + /// @brief BLayout and BDataType are expected to be a tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value, + "BLayout and BDataType must be a tuple."); + + /// @brief CLayout and EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && + !is_detected::value, + "CLayout and EDataType must be a scalar."); + + /// @brief DsLayout and DsDataType are expected to be tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value && + DsLayout::size() == DsDataType::size() && DsLayout::size() > 0, + "DsLayout and DsDataType must be tuples and must have the same size."); + + /// @brief The sizes of NumATensor, NumBTensor and NumDTensor is set by the user." + static constexpr index_t NumATensor = AsDataType::size(); + static constexpr index_t NumBTensor = BsDataType::size(); + static constexpr index_t NumDTensor = DsDataType::size(); + + CK_TILE_HOST static auto GetName() -> const std::string + { + return UniversalGemmKernel::GetName(); + } + + CK_TILE_HOST static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) -> dim3 + { + return UniversalGemmKernel::GridSize(M, N, KBatch); + } + + CK_TILE_HOST static auto MaxOccupancyGridSize(const stream_config& s) -> dim3 + { + return UniversalGemmKernel::MaxOccupancyGridSize(s); + } + + CK_TILE_HOST static constexpr auto BlockSize() -> dim3 + { + return UniversalGemmKernel::BlockSize(); + } + + CK_TILE_HOST static constexpr auto + MakeKernelArgs(const GemmMultiABDHostArgs& hostArgs) -> + typename UniversalGemmKernel::KernelArgs + { + /// @brief Universal GEMM requires array objects and corresponding stride information for + /// matrices A, B, and D. + return UniversalGemmKernel::MakeKernelArgs( + UniversalGemmHostArgs(hostArgs.as_ptr, + hostArgs.bs_ptr, + hostArgs.ds_ptr, + hostArgs.e_ptr, + hostArgs.k_batch, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_As, + hostArgs.stride_Bs, + hostArgs.stride_Ds, + hostArgs.stride_E)); + } + + CK_TILE_HOST static auto + IsSupportedArgument(const typename UniversalGemmKernel::KernelArgs& kargs) -> bool + { + // Currently MultiABD kernel doesn't support k_batch > 1 + if(kargs.k_batch > 1) + { + return false; + } + + return UniversalGemmKernel::IsSupportedArgument(kargs); + } + + CK_TILE_DEVICE auto operator()(typename UniversalGemmKernel::KernelArgs kargs) const -> void + { + UniversalGemmKernel{}.template operator()(kargs); + } +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp index 9d3ac8b901..b0b2905cb4 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp @@ -95,7 +95,7 @@ struct GemmKernelMultiD /// @brief Specify the layout configurations for A, B, E and D using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; using DsLayout = remove_cvref_t; /// @brief Specify the data type configurations for A, B, E and D @@ -114,10 +114,10 @@ struct GemmKernelMultiD !is_detected::value, "BLayout and BDataType must be scalars."); - /// @brief ELayout and EDataType are expected to be scalars, not a tuple. - static_assert(!is_detected::value && + /// @brief CLayout and EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && !is_detected::value, - "ELayout and EDataType must be scalars."); + "CLayout and EDataType must be scalars."); /// @brief DsLayout and DsDataType are expected to be tuple, not a scalar. static_assert(is_detected::value && diff --git a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp index e38e49f5d1..df1d6c9e4f 100644 --- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp @@ -120,10 +120,10 @@ struct GroupedGemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); using OffsetTile1DPartitioner = OffsettedTile1DPartitioner; using Kernel = GroupedGemmKernel; @@ -364,12 +364,8 @@ struct GroupedGemmKernel const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop); // Run GEMM pipeline - const auto& c_block_tile = GemmPipeline{}.template operator()(a_block_window[Base::I0], - b_block_window[Base::I0], - num_loop, - has_hot_loop, - tail_num, - smem_ptr_0); + const auto& c_block_tile = GemmPipeline{}.template operator()( + a_block_window, b_block_window, num_loop, has_hot_loop, tail_num, smem_ptr_0); // Run Epilogue Pipeline auto& c_block_window = gemm_tile_windows.at(Base::I3); EpiloguePipeline{}.template diff --git a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp index cfba8b6c9d..8f44108cc4 100644 --- a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -157,23 +157,23 @@ struct UniversalGemmKernel using EpiloguePipeline = remove_cvref_t; static constexpr bool ADataTypeIsTuple = - is_detected::value; + is_detected::value; static constexpr bool BDataTypeIsTuple = - is_detected::value; + is_detected::value; static constexpr bool DDataTypeIsTuple = is_detected::value; static constexpr bool ALayoutIsTuple = - is_detected::value; + is_detected::value; static constexpr bool BLayoutIsTuple = - is_detected::value; + is_detected::value; static constexpr bool DLayoutIsTuple = is_detected::value; using AsLayout = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using BsLayout = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using DsLayout = std::conditional_t>>; using AsDataType = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using BsDataType = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using DsDataType = @@ -193,9 +193,12 @@ struct UniversalGemmKernel remove_cvref_t, remove_cvref_t>>; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; using EDataType = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + static constexpr index_t kBlockSize = GemmPipeline::BlockSize; // Get the persistent kernel if the pipeline has it available @@ -483,7 +486,7 @@ struct UniversalGemmKernel bool DTesnorIsValid = {true}; static_for<0, NumDTensor, 1>{}([&](auto index) { using DiLayout = remove_cvref_t>; - if(std::is_same_v == false) + if(std::is_same_v == false) { DTesnorIsValid = false; } @@ -529,7 +532,7 @@ struct UniversalGemmKernel } }); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) { @@ -724,7 +727,7 @@ struct UniversalGemmKernel // TODO: enable vector write for C in ColMajor const auto& e_tensor_view = [&]() { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { return make_naive_tensor_view( e_ptr, @@ -818,7 +821,7 @@ struct UniversalGemmKernel // TODO vector write in for C in ColMajor const auto& e_pad_view = [&]() { const auto& e_tensor_view = views.at(I3); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { return pad_tensor_view(e_tensor_view, make_tuple(number{}, @@ -975,8 +978,8 @@ struct UniversalGemmKernel const auto& bs_block_window = gemm_tile_windows.at(I1); const auto& ds_block_window = gemm_tile_windows.at(I2); - const auto& c_block_tile = - GemmPipeline{}(as_block_window[I0], bs_block_window[I0], num_loop, smem_ptr_0); + const auto& c_block_tile = GemmPipeline{}.template operator()( + as_block_window, AElementWise{}, bs_block_window, BElementWise{}, num_loop, smem_ptr_0); if(UseDefaultScheduler || (get_warp_id() == 0)) { @@ -1031,8 +1034,13 @@ struct UniversalGemmKernel const auto& bs_block_window = gemm_tile_windows.at(I1); const auto& ds_block_window = gemm_tile_windows.at(I2); - const auto& c_block_tile = GemmPipeline{}( - as_block_window[I0], bs_block_window[I0], num_loop, smem_ptr_0, smem_ptr_1); + const auto& c_block_tile = GemmPipeline{}.template operator()(as_block_window, + AElementWise{}, + bs_block_window, + BElementWise{}, + num_loop, + smem_ptr_0, + smem_ptr_1); // Run Epilogue Pipeline auto& c_block_window = gemm_tile_windows.at(I3); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp index 2bee550b3c..b5584f98df 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp @@ -11,12 +11,17 @@ namespace ck_tile { template struct GemmPipelineAgBgCrImplBase { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using ADataType = remove_cvref_t{}, AsDataType>>; + using ALayout = remove_cvref_t{}, AsLayout>>; + using BDataType = remove_cvref_t{}, BsDataType>>; + using BLayout = remove_cvref_t{}, BsLayout>>; + static constexpr index_t MPerBlock = BlockGemmShape::kM; static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; @@ -57,6 +62,13 @@ struct GemmPipelineAgBgCrImplBase store_tile(lds_tile_window, block_tile_tmp); } + template + CK_TILE_DEVICE void LocalPrefill(DstTileWindow& lds_tile_window, + const SrcBlockTile& src_block_tile) const + { + store_tile(lds_tile_window, src_block_tile); + } + template CK_TILE_DEVICE void LocalPrefetch(DstBlockTile& dst_block_tile, const SrcTileWindow& lds_tile_window, @@ -88,23 +100,100 @@ struct GemmPipelineAgBgCrImplBase return make_tuple(std::move(a_lds_block), std::move(b_lds_block)); } + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyADramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_col_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp[number{}].get_window_origin() + offset, + Policy::template MakeADramTileDistribution()); + }, + number{}); + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyADramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_col_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = + make_tile_window(dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp.get_window_origin() + offset, + Policy::template MakeADramTileDistribution()); + + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyBDramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_row_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp[number{}].get_window_origin() + offset, + Policy::template MakeBDramTileDistribution()); + }, + number{}); + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyBDramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_row_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = + make_tile_window(dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp.get_window_origin() + offset, + Policy::template MakeBDramTileDistribution()); + + return std::move(a_copy_dram_window); + } + template CK_TILE_DEVICE constexpr auto GetAWindows(const ADramBlockWindowTmp& a_dram_block_window_tmp, const ALdsTensorView& a_lds_block_view, const ALdsLoadTileDistr&, const array& offset = {0, 0}) const { - constexpr bool is_col_major = std::is_same_v; - - using YPerTile = std::conditional_t, number>; - using XPerTile = std::conditional_t, number>; - // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(YPerTile{}, XPerTile{}), - a_dram_block_window_tmp.get_window_origin() + offset, - Policy::template MakeADramTileDistribution()); + auto a_copy_dram_window = CopyADramWindow(a_dram_block_window_tmp, offset); // A LDS tile window for store auto a_lds_shape = []() { @@ -138,16 +227,8 @@ struct GemmPipelineAgBgCrImplBase const BLdsLoadTileDistr&, const array& offset = {0, 0}) const { - constexpr bool is_row_major = std::is_same_v; - - using YPerTile = std::conditional_t, number>; - using XPerTile = std::conditional_t, number>; - - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(YPerTile{}, XPerTile{}), - b_dram_block_window_tmp.get_window_origin() + offset, - Policy::template MakeBDramTileDistribution()); + // A DRAM tile window for load + auto b_copy_dram_window = CopyBDramWindow(b_dram_block_window_tmp, offset); // TODO: Do we really need those two tile windows??? // They're exactly same... diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp index 5f4ee8987e..7159eda683 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp @@ -107,14 +107,23 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 using Base = BaseGemmPipelineAgBgCrCompV3; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; using I0 = number<0>; @@ -386,17 +395,25 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); - - using ABlockTile = - decltype(make_static_distributed_tensor(ABlockTileDistr{})); - using BBlockTile = - decltype(make_static_distributed_tensor(BBlockTileDistr{})); - - ABlockTile a_block_tile; - BBlockTile b_block_tile; - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; @@ -470,45 +476,61 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 // ----------------------------------------------------------------------------------------- // Gemm pipeline start - - // prefetch - // global read 0 - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); - // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); + // LDS write 0 if constexpr(is_a_col_major && !is_a_load_tr_v()) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + // global read 1 + + elementwise_As_res = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); __builtin_amdgcn_sched_barrier(0); @@ -520,38 +542,42 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 { block_sync_lds(); - if constexpr(is_a_col_major && !is_a_load_tr_v()) + if constexpr(is_a_col_major) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } - if constexpr(is_b_row_major && !is_b_load_tr_v()) + if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); HotLoopScheduler(); __builtin_amdgcn_sched_barrier(0); @@ -574,27 +600,26 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); } // __builtin_amdgcn_sched_barrier(0); @@ -602,13 +627,16 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const @@ -628,9 +656,13 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 * @note This is used by the persistent gemm kernel variants that don't determine * hot loop and tail number on the host side, e.g. grouped gemm kernel. */ - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -639,7 +671,7 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -658,20 +690,97 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 * @note This is used by the kernel variants that are able to determine * hot loop and tail number on the host side, e.g. non-persistent gemm kernel. */ - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem); + } + + /** + * @brief Quant operator(), single input: This function runs the pipeline by wrapping it with + * the tail handler. + * + * @note This is used by the persistent gemm kernel variants that don't determine + * hot loop and tail number on the host side, e.g. grouped gemm kernel. + */ + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem); + } + + /** + * @brief Quant operator(), single input: This function runs the pipeline using compile-time + * known hot loop and tail number. + * @param num_loop The number of loop iterations. This is determined at runtime due to e.g. + * SplitK. + * @note This is used by the kernel variants that are able to determine + * hot loop and tail number on the host side, e.g. non-persistent gemm kernel. + */ + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp index c835809b5d..b362f751c6 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp @@ -97,11 +97,24 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 using Base = BaseGemmPipelineAgBgCrCompV4; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; + static_assert(!std::is_same_v, "Not implemented"); static constexpr index_t APackedSize = @@ -109,10 +122,6 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 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 BlockGemm = remove_cvref_t())>; using I0 = number<0>; using I1 = number<1>; @@ -244,18 +253,26 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* __restrict__ p_smem_0, void* __restrict__ p_smem_1) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]), "B block window has incorrect lengths for defined BLayout!"); - ////////////// global window & register ///////////////// - // A DRAM tile window for load - 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(), - Policy::template MakeADramTileDistribution()); - - // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); - - // A register tile for global load - constexpr auto ABlockTileDistr = a_copy_dram_window.get_tile_distribution(); - constexpr auto BBlockTileDistr = b_copy_dram_window.get_tile_distribution(); - using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr)); - using BBlockTile = decltype(make_static_distributed_tensor(BBlockTileDistr)); - ABlockTile a_global_load_tile; - BBlockTile b_global_load_tile; - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; @@ -312,8 +306,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 // global prefetch 0 // global read 0 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + ////////////// LDS desc, window & register ///////////////// auto&& [a_lds_block0, b_lds_block0] = Base::GetABLdsTensorViews(p_smem_0); auto&& [a_lds_block1, b_lds_block1] = Base::GetABLdsTensorViews(p_smem_1); @@ -343,34 +336,75 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + // Generating a tuple with tile_windows for values A0, A1, ... AN + auto a_tile_windows = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); + + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + // Generating a tuple with tile_windows for values B0, B1, ... BN + auto b_tile_windows = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_tile_windows, b_dram_tile_window_step); + // LDS write 0 if constexpr(is_a_col_major && !is_a_load_tr_v()) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } // global read 1 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); block_sync_lds(); constexpr auto ALdsTileDistr = @@ -423,27 +457,32 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window1, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window1, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window1, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window1, elementwise_Bs_res); } - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); if(HasHotLoop) { @@ -461,31 +500,32 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill( - a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill( - b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); // gemm block_gemm(c_block_tile, a_block_tile0, b_block_tile0); HotLoopScheduler(); @@ -501,32 +541,34 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp); } else { - Base::LocalPrefill( - a_copy_lds_window1, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window1, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp); } else { - Base::LocalPrefill( - b_copy_lds_window1, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window1, elementwise_Bs_res); } block_sync_lds(); - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); + // gemm block_gemm(c_block_tile, a_block_tile1, b_block_tile1); HotLoopScheduler(); @@ -548,23 +590,23 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } block_gemm(c_block_tile, a_block_tile0, b_block_tile0); } @@ -606,13 +648,17 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem_0, @@ -628,27 +674,34 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 p_smem_1); } - public: - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, void* __restrict__ p_smem_0, void* __restrict__ p_smem_1) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem_0, p_smem_1); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -658,7 +711,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -670,5 +723,69 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 }; return Base::TailHandler(RunPipeline, has_hot_loop, tail_number); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_0, + void* p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem_0, + p_smem_1); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const index_t num_loop, + void* __restrict__ p_smem_0, + void* __restrict__ p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem_0, + p_smem_1); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* __restrict__ p_smem_0, + void* __restrict__ p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem_0, + p_smem_1); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp index b83d37a790..474d1a5a21 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp @@ -41,15 +41,24 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 using Base = BaseGemmPipelineAgBgCrCompV5; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; @@ -121,17 +130,25 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BsDramBlockWindowTmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* __restrict__ p_smem_0) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v BGemmTile b_tile_0, b_tile_1; // Register tile for A and B. - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); using BBlockTile = decltype(make_static_distributed_tensor(BBlockTileDistr{})); - ABlockTile a_global_load_tile; - BBlockTile b_global_load_tile; + ABlockTile elementwise_As_res; + BBlockTile elementwise_Bs_res; // Block GEMM auto block_gemm = BlockGemm(); @@ -248,33 +267,45 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 // define ping, pong steps here as lambda functions. auto MemoryOpsStep = [&](auto idx) { // Memory read half here. - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each + // A0, A1, … AN. The values A0, A1, … AN are read by the same thread. + elementwise_As_res = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a + // tuple as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each + // B0, B1, … BN. The values B0, B1, … BN are read by the same thread. + elementwise_Bs_res = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a + // tuple as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); if constexpr(is_a_col_major) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } if(idx == 0) @@ -351,13 +382,17 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem_0) const @@ -371,21 +406,62 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 p_smem_0); } - public: - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, void* __restrict__ p_smem_0) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem_0); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_0) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem_0); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const index_t num_loop, + void* __restrict__ p_smem_0) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem_0); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp index e1acfebc47..9e522d4364 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp @@ -157,14 +157,23 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem using Base = BaseGemmPipelineAgBgCrMem; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; @@ -236,17 +245,25 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); @@ -334,10 +353,21 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem // prefetch // global read 0 - Base::GlobalPrefetch( - a_block_tiles.get(I0{}), a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_block_tiles.get(I0{}), b_copy_dram_window, b_dram_tile_window_step); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + a_block_tiles.at(I0{}) = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + b_block_tiles.at(I0{}) = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -348,32 +378,35 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{})); } // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); // main body @@ -397,14 +430,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( a_shuffle_tmp, a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill( a_copy_lds_window, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - a_element_func); + a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { @@ -413,22 +445,23 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( b_shuffle_tmp, b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill( b_copy_lds_window, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - b_element_func); + b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); i += PrefetchStages; @@ -450,26 +483,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(number{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill(a_copy_lds_window, - a_block_tiles.get(number{}), - a_element_func); + a_block_tiles.get(number{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(number{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill(b_copy_lds_window, - b_block_tiles.get(number{}), - b_element_func); + b_block_tiles.get(number{})); } }); @@ -526,17 +557,25 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); @@ -623,10 +664,22 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem // prefetch // global read 0 - Base::GlobalPrefetch( - a_block_tiles.get(I0{}), a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_block_tiles.get(I0{}), b_copy_dram_window, b_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + a_block_tiles.at(I0{}) = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + b_block_tiles.at(I0{}) = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -637,32 +690,35 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{})); } // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); // main body @@ -687,14 +743,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( a_shuffle_tmp, a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill( a_copy_lds_window, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - a_element_func); + a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { @@ -703,22 +758,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( b_shuffle_tmp, b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill( b_copy_lds_window, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - b_element_func); + b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); i += PrefetchStages; @@ -740,26 +797,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(number{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill(a_copy_lds_window, - a_block_tiles.get(number{}), - a_element_func); + a_block_tiles.get(number{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(number{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill(b_copy_lds_window, - b_block_tiles.get(number{}), - b_element_func); + b_block_tiles.get(number{})); } }); @@ -813,13 +868,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const @@ -833,9 +891,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem p_smem); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -844,7 +906,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -856,20 +918,82 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem return Base::TailHandler(RunPipeline, has_hot_loop, tail_number); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const ADataType& a) { e = a; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp index e3b4863392..eb363d59b8 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp @@ -15,14 +15,23 @@ namespace ck_tile { template struct GemmPipelineAGmemBGmemCRegV1 { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; @@ -81,17 +90,25 @@ struct GemmPipelineAGmemBGmemCRegV1 return Policy::template GetSmemSize(); } - template - CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_HOST_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v>, @@ -133,22 +150,30 @@ struct GemmPipelineAGmemBGmemCRegV1 auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); // A DRAM tile window for load - 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(), - Policy::template MakeADramTileDistribution()); + auto as_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); // A LDS tile window for store auto a_copy_lds_window = make_tile_window( a_lds_block, make_tuple(number{}, number{}), {0, 0}); // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); + auto bs_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); // B LDS tile window for store auto b_copy_lds_window = make_tile_window( @@ -182,13 +207,22 @@ struct GemmPipelineAGmemBGmemCRegV1 // prefetch // global read 0 - auto a_block_tile = load_tile(a_copy_dram_window); - auto b_block_tile = load_tile(b_copy_dram_window); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); { // move to 1 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -198,13 +232,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(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); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + store_tile(a_copy_lds_window, a_shuffle_tmp); } else { - store_tile(a_copy_lds_window, tile_elementwise_in(a_element_func, a_block_tile)); + store_tile(a_copy_lds_window, elementwise_As_res); } // LDS write 0 @@ -212,13 +245,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_shuffle_tmp); - store_tile(b_copy_lds_window, b_block_tile_tmp); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + store_tile(b_copy_lds_window, b_shuffle_tmp); } else { - store_tile(b_copy_lds_window, tile_elementwise_in(b_element_func, b_block_tile)); + store_tile(b_copy_lds_window, elementwise_Bs_res); } } @@ -226,8 +258,8 @@ struct GemmPipelineAGmemBGmemCRegV1 while(iCounter > 0) { // global read i + 1 - a_block_tile = load_tile(a_copy_dram_window); - b_block_tile = load_tile(b_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); block_sync_lds(); @@ -237,22 +269,20 @@ struct GemmPipelineAGmemBGmemCRegV1 block_sync_lds(); // move to i + 2 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // LDS write i + 1 if constexpr(is_a_col_major) { auto a_shuffle_tmp_loop = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp_loop, a_block_tile); - store_tile(a_copy_lds_window, - tile_elementwise_in(a_element_func, a_shuffle_tmp_loop)); + transpose_tile2d(a_shuffle_tmp_loop, elementwise_As_res); + store_tile(a_copy_lds_window, a_shuffle_tmp_loop); } else { - 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); + store_tile(a_copy_lds_window, elementwise_As_res); } // LDS write i + 1 @@ -260,14 +290,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto b_shuffle_tmp_loop = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp_loop, b_block_tile); - store_tile(b_copy_lds_window, - tile_elementwise_in(b_element_func, b_shuffle_tmp_loop)); + transpose_tile2d(b_shuffle_tmp_loop, elementwise_Bs_res); + store_tile(b_copy_lds_window, b_shuffle_tmp_loop); } else { - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); } iCounter--; @@ -284,20 +312,40 @@ struct GemmPipelineAGmemBGmemCRegV1 return c_block_tile; } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType & b) { return b; }, + [](auto& e, const BDataType & b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp index b151cd6782..c309f8908a 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp @@ -15,30 +15,66 @@ namespace ck_tile { template struct GemmPipelineAGmemBGmemCRegV2 { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; + static constexpr index_t APackedSize = ck_tile::numeric_traits>::PackedSize; static constexpr index_t BPackedSize = ck_tile::numeric_traits>::PackedSize; - static constexpr index_t kBlockSize = Problem::kBlockSize; + 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; + template + static constexpr index_t GetVectorSizeA() + { + return Problem::VectorSizeA; + } + template + static constexpr index_t GetVectorSizeB() + { + return Problem::VectorSizeB; + } + static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; } + static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool kPadK = Problem::kPadK; + + static constexpr bool Preshuffle = Problem::Preshuffle; + + static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + + // For the basic gemm pipelien DoubleSmemBuffer set to be false naturally. + static constexpr bool DoubleSmemBuffer = false; + [[nodiscard]] CK_TILE_HOST static const std::string GetName() { // clang-format off return concat('_', "pipeline_AGmemBGmemCRegV2", - concat('x', kMPerBlock, kNPerBlock, kKPerBlock, kBlockSize)); + concat('x', kMPerBlock, kNPerBlock, kKPerBlock, BlockSize)); // clang-format on } CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; } @@ -56,17 +92,31 @@ struct GemmPipelineAGmemBGmemCRegV2 BPackedSize; } - template (); + } + + template - CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_HOST_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v>, @@ -98,32 +148,40 @@ struct GemmPipelineAGmemBGmemCRegV2 auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); // A DRAM tile window for load - 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(), - Policy::template MakeADramTileDistribution()); + auto as_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); // A LDS tile window for store auto a_copy_lds_window = make_tile_window(a_lds_block, make_tuple(number{}, number{}), {0, 0}, - a_copy_dram_window.get_tile_distribution()); + as_copy_dram_window[number<0>{}].get_tile_distribution()); // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); + auto bs_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); // B LDS tile window for store auto b_copy_lds_window = make_tile_window(b_lds_block, make_tuple(number{}, number{}), {0, 0}, - b_copy_dram_window.get_tile_distribution()); + bs_copy_dram_window[number<0>{}].get_tile_distribution()); // Block GEMM constexpr auto block_gemm = Policy::template GetBlockGemm(); @@ -153,28 +211,30 @@ struct GemmPipelineAGmemBGmemCRegV2 // prefetch // global read 0 - auto a_block_tile = load_tile(a_copy_dram_window); - auto b_block_tile = load_tile(b_copy_dram_window); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); { // move to 1 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); // LDS write 0 - 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); + store_tile(a_copy_lds_window, elementwise_As_res); // global read 1 - a_block_tile = load_tile(a_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); // LDS write 0 - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); // global read 1 - b_block_tile = load_tile(b_copy_dram_window); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); } index_t iCounter = num_loop - 2; @@ -189,20 +249,18 @@ struct GemmPipelineAGmemBGmemCRegV2 block_sync_lds(); // move to i + 2 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_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); + store_tile(a_copy_lds_window, elementwise_As_res); // global read i + 2 - a_block_tile = load_tile(a_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); // LDS write i + 1 - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); // global read i + 2 - b_block_tile = load_tile(b_copy_dram_window); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); iCounter--; @@ -218,11 +276,9 @@ struct GemmPipelineAGmemBGmemCRegV2 block_sync_lds(); // LDS write num_loop - 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); + store_tile(a_copy_lds_window, elementwise_As_res); - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); block_sync_lds(); @@ -241,12 +297,28 @@ struct GemmPipelineAGmemBGmemCRegV2 { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType & b) { return b; }, + [](auto& e, const BDataType & b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp index 52bd07c9e2..c73fa29245 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp @@ -5,16 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { -template @@ -22,18 +25,49 @@ struct GemmPipelineProblemBase { using Traits = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; // actually AccDataType - using ComputeDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; // actually AccDataType static constexpr bool FixedVectorSize = FixedVectorSize_; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + static constexpr bool ComputeDataTypeIsTuple = is_detected::value; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + static constexpr bool ALayoutIsTuple = is_detected::value; + static constexpr bool BLayoutIsTuple = is_detected::value; + + using ComputeDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + using AsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + using BsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ComputeDataType = remove_cvref_t{}, ComputeDataTypeTuple>>; + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using ALayout = remove_cvref_t{}, AsLayoutTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using BLayout = remove_cvref_t{}, BsLayoutTuple>>; static constexpr bool TransposeC = Traits::TransposeC; static constexpr index_t NumWaveGroups = Traits::NumWaveGroups; @@ -66,7 +100,7 @@ struct GemmPipelineProblemBase { constexpr index_t PackedSize = ck_tile::numeric_traits>::PackedSize; - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { constexpr index_t pixels_per_thread = BlockGemmShape::kM * BlockGemmShape::kK / kBlockSize; @@ -84,7 +118,7 @@ struct GemmPipelineProblemBase { constexpr index_t PackedSize = ck_tile::numeric_traits>::PackedSize; - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { constexpr index_t pixels_per_thread = BlockGemmShape::kN * BlockGemmShape::kK / kBlockSize; @@ -125,7 +159,7 @@ struct GemmPipelineProblemBase { return VectorSizeA_; } - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) { return kPadK ? 1 : GetAlignmentA(); } @@ -140,7 +174,7 @@ struct GemmPipelineProblemBase { return VectorSizeB_; } - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) { return kPadN ? 1 : GetAlignmentB(); } @@ -161,35 +195,40 @@ struct GemmPipelineProblemBase }(); }; -// Alias for GemmPipelineProblem -template -using GemmPipelineProblem = GemmPipelineProblemBase; -template @@ -197,18 +236,48 @@ struct UniversalGemmPipelineProblem { using Traits = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; // actually AccDataType - using ComputeDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; // actually AccDataType + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; static constexpr bool FixedVectorSize = FixedVectorSize_; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + static constexpr bool ComputeDataTypeIsTuple = is_detected::value; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + static constexpr bool ALayoutIsTuple = is_detected::value; + static constexpr bool BLayoutIsTuple = is_detected::value; + + using ComputeDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + using AsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + using BsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ComputeDataType = remove_cvref_t{}, ComputeDataTypeTuple>>; + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using ALayout = remove_cvref_t{}, AsLayoutTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using BLayout = remove_cvref_t{}, BsLayoutTuple>>; static constexpr bool TransposeC = Traits::TransposeC; static constexpr index_t NumWaveGroups = Traits::NumWaveGroups; diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp index 8d47ab878e..c8f874acd6 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -356,11 +356,14 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeA() { - using ALayout = remove_cvref_t; - using ADataType = remove_cvref_t; + using AsLayout = remove_cvref_t; + using AsDataType = remove_cvref_t; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + using ALayout = remove_cvref_t{}, AsLayout>>; + using ADataType = remove_cvref_t{}, AsDataType>>; + if constexpr(std::is_same_v) { return GetGlobalVectorLoadSize CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeB() { - using BLayout = remove_cvref_t; - using BDataType = remove_cvref_t; + using BsLayout = remove_cvref_t; + using BsDataType = remove_cvref_t; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + using BLayout = remove_cvref_t{}, BsLayout>>; + using BDataType = remove_cvref_t{}, BsDataType>>; + if constexpr(std::is_same_v) { return GetGlobalVectorLoadSize CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { - using ALayout = remove_cvref_t; - constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; @@ -491,6 +495,8 @@ struct UniversalGemmBasePolicy Problem::FixedVectorSize ? Problem::VectorSizeA : GetVectorSizeA(); constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + using ALayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; // Tile: MPerBlock X KPerBlock if constexpr(std::is_same_v) { @@ -518,8 +524,6 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeBDramTileDistribution() { - using BLayout = remove_cvref_t; - constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; @@ -527,6 +531,8 @@ struct UniversalGemmBasePolicy Problem::FixedVectorSize ? Problem::VectorSizeB : GetVectorSizeB(); constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + using BLayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; // Tile: KPerBlock X NPerBlock if constexpr(std::is_same_v) { @@ -554,7 +560,8 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledARegTileDistribution() { - using ALayout = remove_cvref_t; + using ALayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; static_assert(std::is_same_v); constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; @@ -574,7 +581,8 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledBRegTileDistribution() { - using BLayout = remove_cvref_t; + using BLayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; static_assert(std::is_same_v); constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; diff --git a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp index 64900c9a97..96203b2cd2 100644 --- a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp +++ b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp @@ -10,8 +10,8 @@ namespace ck_tile { template struct TileGemmTraits @@ -23,9 +23,9 @@ struct TileGemmTraits // TODO this can't be hardcoded here! Should be in policy! static constexpr int _VectorSize = 16; - using ALayout = ALayout_; - using BLayout = BLayout_; - using CLayout = CLayout_; + using AsLayout = AsLayout_; + using BsLayout = BsLayout_; + using CLayout = CLayout_; static constexpr bool TransposeC = false; static constexpr bool UseStructuredSparsity = false; @@ -36,8 +36,8 @@ template @@ -76,8 +76,8 @@ using PersistentTileGemmUniversalTraits = TileGemmUniversalTraits { - using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV1; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV1; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockWeightPreshuffle = remove_cvref_t())>; @@ -188,7 +197,12 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 } } - template + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -455,7 +469,33 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 return c_block_tile; } - template + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + [[maybe_unused]] const AElementFunction& a_element_func, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + [[maybe_unused]] const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()( + a_dram_block_window_tmp[number<0>{}], + [](const ADataType & a) { return a; }, + b_flat_dram_block_window_tmp[number<0>{}], + num_loop, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, @@ -463,7 +503,7 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_flat_dram_block_window_tmp, num_loop, p_smem); diff --git a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp index 129eac6557..356ad91448 100644 --- a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp @@ -53,14 +53,23 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 { using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV2; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockWeightPreshuffle = remove_cvref_t())>; @@ -502,7 +511,10 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 template + typename AElementFunction, + typename std::enable_if_t::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -1001,8 +1013,37 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 return c_block_tile; } + // called from universal gemm kernel + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + [[maybe_unused]] const AElementFunction& a_element_func, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + [[maybe_unused]] const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_ping, + void* p_smem_pong) const + { + return operator()( + a_dram_block_window_tmp[number<0>{}], + [](const ADataType& a) { return a; }, + b_flat_dram_block_window_tmp[number<0>{}], + num_loop, + p_smem_ping, + p_smem_pong); + } + // called from general gemm kernel - template + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, @@ -1019,9 +1060,13 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 } // called from grouped gemm kernel - template + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_flat_dram_block_window_tmp, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, TailNumber tail_number, void* __restrict__ p_smem_0, diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp b/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp index 44c6cd66c6..f505efe4e0 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp +++ b/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp @@ -44,6 +44,10 @@ struct TileGemmQuantTraits using AQLayout = AQLayout_; using BQLayout = BQLayout_; + // TODO: It should be replaced to single value + using AsLayout = ALayout_; + using BsLayout = BLayout_; + static constexpr bool TransposeC = false; static constexpr bool UseStructuredSparsity = false; static constexpr index_t NumWaveGroups = 1; diff --git a/test/ck_tile/CMakeLists.txt b/test/ck_tile/CMakeLists.txt index 9314d4b795..b08f0d8316 100644 --- a/test/ck_tile/CMakeLists.txt +++ b/test/ck_tile/CMakeLists.txt @@ -5,6 +5,7 @@ add_subdirectory(batched_gemm) add_subdirectory(grouped_gemm) add_subdirectory(grouped_gemm_preshuffle) add_subdirectory(gemm_multi_d) +add_subdirectory(gemm_multi_abd) add_subdirectory(gemm_streamk) add_subdirectory(data_type) add_subdirectory(container) diff --git a/test/ck_tile/gemm_multi_abd/CMakeLists.txt b/test/ck_tile/gemm_multi_abd/CMakeLists.txt new file mode 100644 index 0000000000..ac3b59d5d3 --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/CMakeLists.txt @@ -0,0 +1,12 @@ +# Currently ck_tile is only built on gfx9 +set(EXAMPLE_GEMM_COMPILE_OPTIONS) +if(CK_USE_OCP_FP8) + list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8) +endif() + +if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95") + add_gtest_executable(test_gemm_multi_abd_cshuffle test_gemm_multi_abd_cshuffle.cpp) + add_gtest_executable(test_gemm_multi_abd_default2d test_gemm_multi_abd_default2d.cpp) + target_compile_definitions(test_gemm_multi_abd_cshuffle PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + target_compile_definitions(test_gemm_multi_abd_default2d PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) +endif() diff --git a/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_cshuffle.cpp b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_cshuffle.cpp new file mode 100644 index 0000000000..9821963458 --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_cshuffle.cpp @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include "gtest/gtest.h" + +#include "ck_tile/host.hpp" +#include "test_gemm_multi_abd_util.hpp" + +using F16 = ck_tile::half_t; +using BF16 = ck_tile::bf16_t; +using F32 = float; +using F8 = ck_tile::fp8_t; + +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + +// clang-format off +using KernelTypes = ::testing::Types< + // Has cshuffle epilogue enabled + // A0Layout, A1Layout, B0Layout, B1Layout CLayout, D0Layout, D1Layout, A0DataType, A01DataType B0DataType, B0DataType, D0DataType, D1DataType, AccDataType, EDataType, AElementWiseFn, BElementWiseFn, CDElementWiseFn, UseCshuffleEpilog + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, BF16, BF16, F32, F16, AddScale, AddScale, ElementWiseAddAdd, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, ElementWiseAddAdd, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, ElementWiseAddAdd, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, F32, AddScale, AddScale, ElementWiseAddAdd, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, F8, F8, F32, F32, AddScale, AddScale, ElementWiseAddAdd, std::true_type>, + + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F16, F16, F32, F16, AddScale, AddScale, MultiplyMultiply, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, BF16, BF16, F32, F32, AddScale, AddScale, MultiplyMultiply, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F32, AddScale, AddScale, MultiplyMultiply, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, MultiplyMultiply, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, F32, AddScale, AddScale, MultiplyMultiply, std::true_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, F8, F8, F32, F32, AddScale, AddScale, MultiplyMultiply, std::true_type> + >; +// clang-format on + +TYPED_TEST_SUITE(TestCkTileGemmMultiABD, KernelTypes); + +#include "test_gemm_multi_abd_ut_cases_cshuffle.inc" diff --git a/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_default2d.cpp b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_default2d.cpp new file mode 100644 index 0000000000..b3a89aba05 --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_default2d.cpp @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include "gtest/gtest.h" + +#include "ck_tile/host.hpp" +#include "test_gemm_multi_abd_util.hpp" + +using F16 = ck_tile::half_t; +using BF16 = ck_tile::bf16_t; +using F32 = float; +using F8 = ck_tile::fp8_t; + +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + +// clang-format off +using KernelTypes = ::testing::Types< + // Has cshuffle epilogue disabled + // A0Layout, A1Layout, B0Layout, B1Layout CLayout, D0Layout, D1Layout, A0DataType, A01DataType B0DataType, B0DataType, D0DataType, D1DataType, AccDataType, EDataType, AElementWiseFn, BElementWiseFn, CDElementWiseFn, UseCshuffleEpilog + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, F32, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F32, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, BF16, BF16, F32, BF16, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, BF16, AddScale, AddScale, ElementWiseAddAdd, std::false_type>, + + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F16, F16, F32, F16, AddScale, AddScale, MultiplyMultiply, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F16, AddScale, AddScale, MultiplyMultiply, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, F32, AddScale, AddScale, MultiplyMultiply, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, F32, F32, F32, F32, AddScale, AddScale, MultiplyMultiply, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F16, F16, F16, F16, BF16, BF16, F32, BF16, AddScale, AddScale, MultiplyMultiply, std::false_type>, + std::tuple< Row, Row, Col, Col, Row, Row, Row, F8, F8, F8, F8, BF16, BF16, F32, BF16, AddScale, AddScale, MultiplyMultiply, std::false_type> + >; +// clang-format on + +TYPED_TEST_SUITE(TestCkTileGemmMultiABD, KernelTypes); + +#include "test_gemm_multi_abd_ut_cases_default2d.inc" diff --git a/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_cshuffle.inc b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_cshuffle.inc new file mode 100644 index 0000000000..5aa113608f --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_cshuffle.inc @@ -0,0 +1,211 @@ +#pragma once + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_256x512x256) +{ + constexpr int M = 256; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_512x256x256) +{ + constexpr int M = 512; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_512x512x256) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_256x256x256) +{ + constexpr int M = 256; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_512x768x256) +{ + constexpr int M = 512; + constexpr int N = 768; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_512x1280x256) +{ + constexpr int M = 512; + constexpr int N = 1280; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_256x1280x256) +{ + constexpr int M = 256; + constexpr int N = 1280; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_768x512x256) +{ + constexpr int M = 768; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_1280x512x256) +{ + constexpr int M = 1280; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1CShuffle_1280x256x256) +{ + constexpr int M = 1280; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_512x512x512) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_256x512x256) +{ + constexpr int M = 256; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_512x256x256) +{ + constexpr int M = 512; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_512x512x256) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_256x256x256) +{ + constexpr int M = 256; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_512x768x256) +{ + constexpr int M = 512; + constexpr int N = 768; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_512x1280x256) +{ + constexpr int M = 512; + constexpr int N = 1280; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_256x1280x256) +{ + constexpr int M = 256; + constexpr int N = 1280; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_768x512x256) +{ + constexpr int M = 768; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_1280x512x256) +{ + constexpr int M = 1280; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2CShuffle_1280x256x256) +{ + constexpr int M = 1280; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} diff --git a/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_default2d.inc b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_default2d.inc new file mode 100644 index 0000000000..cc7603164c --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_ut_cases_default2d.inc @@ -0,0 +1,211 @@ +#pragma once + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_256x512x256) +{ + constexpr int M = 256; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_512x256x256) +{ + constexpr int M = 512; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_512x512x256) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_256x256x256) +{ + constexpr int M = 256; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_512x768x256) +{ + constexpr int M = 512; + constexpr int N = 768; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_512x1280x256) +{ + constexpr int M = 512; + constexpr int N = 1280; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_256x1280x256) +{ + constexpr int M = 256; + constexpr int N = 1280; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_768x512x256) +{ + constexpr int M = 768; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_1280x512x256) +{ + constexpr int M = 1280; + constexpr int N = 512; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch1Default_1280x256x256) +{ + constexpr int M = 1280; + constexpr int N = 256; + constexpr int K = 256; + constexpr int kBatch = 1; + + EXPECT_EQ(this->Run(M, N, K, kBatch), true); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_512x512x512) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_256x512x256) +{ + constexpr int M = 256; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_512x256x256) +{ + constexpr int M = 512; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_512x512x256) +{ + constexpr int M = 512; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_256x256x256) +{ + constexpr int M = 256; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_512x768x256) +{ + constexpr int M = 512; + constexpr int N = 768; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_512x1280x256) +{ + constexpr int M = 512; + constexpr int N = 1280; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_256x1280x256) +{ + constexpr int M = 256; + constexpr int N = 1280; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_768x512x256) +{ + constexpr int M = 768; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_1280x512x256) +{ + constexpr int M = 1280; + constexpr int N = 512; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmMultiABD, TestCkTileGemmMultiABDKBatch2Default_1280x256x256) +{ + constexpr int M = 1280; + constexpr int N = 256; + constexpr int K = 512; + constexpr int kBatch = 2; + + EXPECT_THROW(this->Run(M, N, K, kBatch), std::runtime_error); +} diff --git a/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_util.hpp b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_util.hpp new file mode 100644 index 0000000000..428bed4e25 --- /dev/null +++ b/test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_util.hpp @@ -0,0 +1,500 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +#pragma once + +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" + +struct AddScale +{ + template + CK_TILE_HOST_DEVICE constexpr void operator()(E& a, const A0& a0, const A1& a1) const + { + a = scale * (ck_tile::type_convert(a0) + ck_tile::type_convert(a1)); + } + + float scale = 1.0; +}; + +struct MultiplyMultiply +{ + template + CK_TILE_HOST_DEVICE auto operator()(E& e, const C& c, const D0& d0, const D1& d1) const -> void + { + const float x0_f = ck_tile::type_convert(c) * ck_tile::type_convert(d0) * + ck_tile::type_convert(d1); + + e = ck_tile::type_convert(x0_f); + } +}; + +struct ElementWiseAddAdd +{ + template + CK_TILE_HOST_DEVICE auto operator()(E& e, const C& c, const D0& d0, const D1& d1) const -> void + { + const float x0_f = ck_tile::type_convert(c) + ck_tile::type_convert(d0) + + ck_tile::type_convert(d1); + + e = ck_tile::type_convert(x0_f); + } +}; + +template +static constexpr inline auto is_row_major(Layout layout_) +{ + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +template +auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeTypeAB = + std::conditional_t; + + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} + +template +class TestCkTileGemmMultiABD : public ::testing::Test +{ + protected: + using A0Layout = std::tuple_element_t<0, Tuple>; + using A1Layout = std::tuple_element_t<1, Tuple>; + using B0Layout = std::tuple_element_t<2, Tuple>; + using B1Layout = std::tuple_element_t<3, Tuple>; + using D0Layout = std::tuple_element_t<4, Tuple>; + using D1Layout = std::tuple_element_t<5, Tuple>; + using ELayout = std::tuple_element_t<6, Tuple>; + using A0DataType = std::tuple_element_t<7, Tuple>; + using A1DataType = std::tuple_element_t<8, Tuple>; + using B0DataType = std::tuple_element_t<9, Tuple>; + using B1DataType = std::tuple_element_t<10, Tuple>; + using D0DataType = std::tuple_element_t<11, Tuple>; + using D1DataType = std::tuple_element_t<12, Tuple>; + using AccDataType = std::tuple_element_t<13, Tuple>; + using EDataType = std::tuple_element_t<14, Tuple>; + using AElementWiseFn = std::tuple_element_t<15, Tuple>; + using BElementWiseFn = std::tuple_element_t<16, Tuple>; + using CDElementWiseFn = std::tuple_element_t<17, Tuple>; + using UseCshuffleEpilog = std::tuple_element_t<18, Tuple>; + + using AsLayout = ck_tile::tuple; + using AsDataType = ck_tile::tuple; + using BsLayout = ck_tile::tuple; + using BsDataType = ck_tile::tuple; + using DsLayout = ck_tile::tuple; + using DsDataType = ck_tile::tuple; + + template + void invoke_gemm_multi_abd(const ck_tile::GemmMultiABDHostArgs& args, + const ck_tile::stream_config& s) + { + constexpr ck_tile::index_t M_Tile = 256; + constexpr ck_tile::index_t N_Tile = 256; + constexpr ck_tile::index_t K_Tile = 32; + + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = 32; + constexpr ck_tile::index_t N_Warp_Tile = 32; + constexpr ck_tile::index_t K_Warp_Tile = 16; + + constexpr bool DoubleSmemBuffer = false; + + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + + constexpr bool TransposeC = false; + + constexpr int kBlockPerCu = 1; + constexpr ck_tile::index_t TileParitionerGroupNum = 8; + constexpr ck_tile::index_t TileParitionerM01 = 4; + + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; + + using Traits = ck_tile::TileGemmTraits; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmPipelineProblem = + ck_tile::GemmPipelineProblem; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t k_grain = args.k_batch * K_Tile; + const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + float ave_time{0}; + + const auto Run = [&](const auto has_hot_loop_, + const auto tail_number_, + const auto memory_operation_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr auto scheduler = ck_tile::GemmPipelineScheduler::Intrawave; + constexpr auto memory_operation = memory_operation_.value; + + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + + using DefaultGemmEpilogue = ck_tile::DefaultGemm2DEpilogue< + ck_tile::DefaultGemm2DEpilogueProblem>; + + using CShuffleGemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem>; + + using GemmEpilogue = std:: + conditional_t; + + using Kernel = ck_tile::GemmKernelMultiABD; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' + << "shape: " << GemmShape::GetName() << '\n' + << "problem: " << GemmPipelineProblem::GetName() << '\n' + << "pipeline: " << GemmPipeline::GetName() << '\n' + << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" + << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z + << "}" << std::endl; + } + + ave_time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + return ave_time; + }; + + const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) { + if(args.k_batch == 1) + { + std::cout << "Run without SplitK" << std::endl; + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + else + { + std::cout << "Run using SplitK" << std::endl; + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + }; + + BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num); + } + + public: + bool Run(const int M, + const int N, + const int K, + const int k_batch, + int StrideA0 = 0, + int StrideA1 = 0, + int StrideB0 = 0, + int StrideB1 = 0, + int StrideD0 = 0, + int StrideD1 = 0, + int StrideE = 0) + { + using namespace ck_tile::literals; + + auto f_host_tensor_descriptor = [](std::size_t row, + std::size_t col, + std::size_t stride, + auto layout) { + if constexpr(std::is_same_v) + { + return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride}); + } + }; + + auto f_get_default_stride = + [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { + if(stride == 0) + { + if constexpr(std::is_same_v) + { + return col; + } + else + { + return row; + } + } + else + return stride; + }; + + StrideA0 = f_get_default_stride(M, K, StrideA0, A0Layout{}); + StrideA1 = f_get_default_stride(M, K, StrideA1, A1Layout{}); + + StrideB0 = f_get_default_stride(K, N, StrideB0, B0Layout{}); + StrideB1 = f_get_default_stride(K, N, StrideB1, B1Layout{}); + + StrideD0 = f_get_default_stride(M, N, StrideD0, D0Layout{}); + StrideD1 = f_get_default_stride(M, N, StrideD1, D1Layout{}); + + StrideE = f_get_default_stride(M, N, StrideE, ELayout{}); + + ck_tile::HostTensor a0_m_k_tesnor( + f_host_tensor_descriptor(M, K, StrideA0, A0Layout{})); + ck_tile::HostTensor a1_m_k_tesnor( + f_host_tensor_descriptor(M, K, StrideA1, A1Layout{})); + + ck_tile::HostTensor b0_k_n_tensors( + f_host_tensor_descriptor(K, N, StrideB0, B0Layout{})); + ck_tile::HostTensor b1_k_n_tensors( + f_host_tensor_descriptor(K, N, StrideB1, B1Layout{})); + + ck_tile::HostTensor d0_m_n_tensors( + f_host_tensor_descriptor(M, N, StrideD0, D0Layout{})); + ck_tile::HostTensor d1_m_n_tensors( + f_host_tensor_descriptor(M, N, StrideD1, D1Layout{})); + + ck_tile::HostTensor e_m_n_device_result( + f_host_tensor_descriptor(M, N, StrideE, ELayout{})); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a0_m_k_tesnor); + ck_tile::FillUniformDistribution{-1.f, 1.f}(a1_m_k_tesnor); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(b0_k_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b1_k_n_tensors); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(d0_m_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(d1_m_n_tensors); + + ck_tile::DeviceMem a0_m_k_dev_buf(a0_m_k_tesnor.get_element_space_size_in_bytes()); + ck_tile::DeviceMem a1_m_k_dev_buf(a1_m_k_tesnor.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem b0_k_n_dev_buf(b0_k_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b1_k_n_dev_buf(b1_k_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem d0_m_n_dev_buf(d0_m_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d1_m_n_dev_buf(d1_m_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem e_m_n_dev_buf(e_m_n_device_result.get_element_space_size_in_bytes()); + + a0_m_k_dev_buf.ToDevice(a0_m_k_tesnor.mData.data()); + a1_m_k_dev_buf.ToDevice(a1_m_k_tesnor.mData.data()); + + b0_k_n_dev_buf.ToDevice(b0_k_n_tensors.mData.data()); + b1_k_n_dev_buf.ToDevice(b1_k_n_tensors.mData.data()); + + d0_m_n_dev_buf.ToDevice(d0_m_n_tensors.mData.data()); + d1_m_n_dev_buf.ToDevice(d1_m_n_tensors.mData.data()); + + e_m_n_dev_buf.SetZero(); + e_m_n_device_result.SetZero(); + + std::array as_ptr_buf = {a0_m_k_dev_buf.GetDeviceBuffer(), + a1_m_k_dev_buf.GetDeviceBuffer()}; + + std::array bs_ptr_buf = {b0_k_n_dev_buf.GetDeviceBuffer(), + b1_k_n_dev_buf.GetDeviceBuffer()}; + + std::array ds_ptr_buf = {d0_m_n_dev_buf.GetDeviceBuffer(), + d1_m_n_dev_buf.GetDeviceBuffer()}; + + std::array strideAs = {StrideA0, StrideA1}; + std::array strideBs = {StrideB0, StrideB1}; + std::array strideDs = {StrideD0, StrideD1}; + + ck_tile::GemmMultiABDHostArgs + args({as_ptr_buf, + bs_ptr_buf, + ds_ptr_buf, + e_m_n_dev_buf.GetDeviceBuffer(), + k_batch, + M, + N, + K, + strideAs, + strideBs, + strideDs, + StrideE}); + + invoke_gemm_multi_abd(args, ck_tile::stream_config{nullptr, false}); + + std::cout << "Run kernel with M =" << M << " N =" << N << " K =" << K + << " StrideA0 =" << StrideA0 << " StrideA1 =" << StrideA1 + << " StrideB0 =" << StrideB0 << " StrideB1 =" << StrideB1 + << " StrideE =" << StrideE << " StrideD0 =" << StrideD0 + << " StrideD1 =" << StrideD1 << std::endl; + + e_m_n_dev_buf.FromDevice(e_m_n_device_result.data()); + bool pass = true; + + ck_tile::HostTensor a_m_k_host_ref_element_result( + f_host_tensor_descriptor(M, K, StrideA0, A0Layout{})); + ck_tile::HostTensor b_k_n_host_ref_element_result( + f_host_tensor_descriptor(K, N, StrideB0, B0Layout{})); + ck_tile::HostTensor e_m_n_host_ref( + f_host_tensor_descriptor(M, N, StrideE, ELayout{})); + a_m_k_host_ref_element_result.SetZero(); + b_k_n_host_ref_element_result.SetZero(); + e_m_n_host_ref.SetZero(); + + ck_tile::reference_gemm_multiple_abd({a0_m_k_tesnor, a1_m_k_tesnor}, + {b0_k_n_tensors, b1_k_n_tensors}, + {d0_m_n_tensors, d1_m_n_tensors}, + a_m_k_host_ref_element_result, + b_k_n_host_ref_element_result, + e_m_n_host_ref); + + const float max_accumulated_value = + *std::max_element(e_m_n_host_ref.mData.begin(), e_m_n_host_ref.mData.end()); + const auto rtol_atol = + calculate_rtol_atol( + K, k_batch, max_accumulated_value); + pass = ck_tile::check_err(e_m_n_device_result, + e_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + + return pass; + } +}; From 47cd0d5cff77658adc1c9f184c012ec3496e8214 Mon Sep 17 00:00:00 2001 From: SamiAario-AMD Date: Fri, 19 Sep 2025 07:26:10 +0300 Subject: [PATCH 02/12] Add gemm weight preshuffle pk_int_t support (#2858) * Factor out the three separate copies of load_interleaved_pk_type into a common utility class * Add preprocessing with optional cache flushing and clearing of output for k_batch > 1 to the weight preshuffle GEMM example * Remove a duplicate function * Add support for B tensor type pk_int4_t for the weight preshuffle GEMM, with tests included * I4 support introduced more failing test cases that mirror the existing ones for F8 * Simplify the check for which tests to skip (they all have F8 as A tensor type) * Add a changelog entry * add the test for v2 wp pipeline, polish the code, add the support of int4 for v2 wp pipeline * have a workable version for atomic add * Revert "have a workable version for atomic add" This reverts commit 792377a590c26cfff9c8f545d9a9e8484a7422eb. --------- Co-authored-by: ThomasNing --- CHANGELOG.md | 1 + .../ops/common/load_interleaved_pk_type.hpp | 58 +++++++++++++++++++ .../block/block_universal_gemm_as_bs_cr.hpp | 37 ++++-------- ..._pipeline_agmem_bgmem_creg_base_policy.hpp | 18 +++--- .../wp_pipeline_agmem_bgmem_creg_v1.hpp | 28 +++++---- .../wp_pipeline_agmem_bgmem_creg_v2.hpp | 28 +++++---- .../block_universal_gemm_as_aquant_bs_cr.hpp | 30 +++------- .../block_universal_gemm_as_bs_bquant_cr.hpp | 31 +++------- .../test_batched_gemm_ut_cases.inc | 3 +- .../test_gemm_pipeline_smoke_run_test.inc | 57 +----------------- .../test_gemm_pipeline_kernel_types.hpp | 25 ++++---- .../test_gemm_pipeline_ut_cases.inc | 8 +-- .../test_gemm_pipeline_util.hpp | 36 +++++++++--- 13 files changed, 183 insertions(+), 177 deletions(-) create mode 100644 include/ck_tile/ops/common/load_interleaved_pk_type.hpp diff --git a/CHANGELOG.md b/CHANGELOG.md index dafe1b5c87..6dd06195c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj ## Composable Kernel 1.2.0 for ROCm 7.0.0 ### Added +* Added support for B Tensor type pk_int4_t in the CK TILE weight preshuffle GEMM. * Added support for B Tensor Preshuffle in CK TILE Grouped GEMM. * Added a basic copy kernel example and supporting documentation for new CK Tile developers. * Added support for bf16, f32, and f16 for 2D and 3D NGCHW grouped convolution backward data diff --git a/include/ck_tile/ops/common/load_interleaved_pk_type.hpp b/include/ck_tile/ops/common/load_interleaved_pk_type.hpp new file mode 100644 index 0000000000..f8432b9da0 --- /dev/null +++ b/include/ck_tile/ops/common/load_interleaved_pk_type.hpp @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core/config.hpp" +#include "ck_tile/ops/elementwise.hpp" + +namespace ck_tile { + +template +struct is_pk_int4 : std::false_type +{ +}; +template <> +struct is_pk_int4 : std::true_type +{ +}; + +template +struct InterleavedPKTypeLoader +{ + template + CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, + const WarpWindow& warp_window) + { + const element_wise::PassThroughPack8 elementwise_op{}; + + static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); + constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; + const auto in_dstr_tensors = load_tile(warp_window); + + using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); + static_for<0, thread_buffer_size, 1>{}([&](auto i) { + elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), + in_dstr_tensors.get_thread_buffer().template get_as()[i]); + }); + } +}; + +template +CK_TILE_DEVICE void load_int4_tile(WarpTile& dst, const WarpWindow& src) +{ + if constexpr(is_pk_int4>::value) + { + InterleavedPKTypeLoader::load_interleaved_pk_type(dst, src); + } + else + { + dst = load_tile(src); + } +} + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp index e1b0792ecf..94adb42880 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" @@ -13,7 +14,9 @@ namespace ck_tile { // A is block window on shared memory // B is block window on shared memory // C is block distributed tensor -template +template struct BlockUniversalGemmAsBsCr { private: @@ -91,6 +94,7 @@ struct BlockUniversalGemmAsBsCr using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -179,25 +183,6 @@ struct BlockUniversalGemmAsBsCr return b_block_dstr_encode; } - private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - template struct BlockGemmImpl { @@ -239,7 +224,7 @@ struct BlockUniversalGemmAsBsCr if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -247,7 +232,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { @@ -317,7 +302,7 @@ struct BlockUniversalGemmAsBsCr { if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else if constexpr(ALoadTranspose) { @@ -329,7 +314,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else if constexpr(BLoadTranspose) { @@ -468,7 +453,7 @@ struct BlockUniversalGemmAsBsCr if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else if constexpr(ALoadTranspose) { @@ -480,7 +465,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else if constexpr(BLoadTranspose) { diff --git a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp index 71ca907c07..f1c8f2ec9b 100644 --- a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp @@ -289,13 +289,17 @@ struct UniversalWeightPreshufflePipelineAgBgCrPolicy { using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmDispatcher; + using BTypeToUse = + std::conditional_t, + typename Problem::ADataType, + typename Problem::BDataType>; + using WarpGemm = WarpGemmDispatcher; using BlockWeightPreshufflePolicy = BlockWeightPreshuffleASmemBSmemCRegV1CustomPolicy::value && !is_detected::value, - bool>* = nullptr> + bool>* = nullptr, + index_t UnaryOpSize_ = 8> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -310,14 +312,14 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 NIterPerWarp> b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + using BTypeToUse = + std::conditional_t, ADataType, BDataType>; + using BTileType = decltype(make_static_distributed_tensor(b_flat_distribution)); + + statically_indexed_array, NIterPerWarp> b_warp_tensor; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_2; static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { @@ -327,7 +329,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -375,7 +378,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_2(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_2(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -408,7 +412,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -445,7 +450,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_2(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_2(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); diff --git a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp index 356ad91448..670f4b0575 100644 --- a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/host/concat.hpp" #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp" @@ -514,7 +515,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 typename AElementFunction, typename std::enable_if_t::value && !is_detected::value, - bool>* = nullptr> + bool>* = nullptr, + index_t UnaryOpSize_ = 8> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -631,19 +633,19 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 b_flat_distribution); // pingpong buffer for B + using BTypeToUse = + std::conditional_t, ADataType, BDataType>; + using BTileType = decltype(make_static_distributed_tensor(b_flat_distribution)); + statically_indexed_array< statically_indexed_array, NIterPerWarp> b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_ping; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_pong; // Prefetch A0 @@ -659,7 +661,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_ping(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); // move B window to next flat K @@ -706,7 +709,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_pong(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -782,7 +786,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_ping(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -862,7 +867,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_pong(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); diff --git a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp b/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp index 182d9251b1..f75d02f1a6 100644 --- a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp +++ b/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp @@ -5,19 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" namespace ck_tile { -template +template struct BlockGemmAQuantBase { using AQDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; - static constexpr index_t UnaryOpSize = UnaryOpSize_; template CK_TILE_DEVICE static float cvt_scale_to_fp32(T scale) { @@ -42,23 +42,6 @@ struct BlockGemmAQuantBase } return scale_reg_f; } - - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - const element_wise::PassThroughPack8 elementwise_op{}; - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } }; // A is block window on shared memory @@ -66,7 +49,9 @@ struct BlockGemmAQuantBase // Consecutive kQuantGroupSize elements of A are quantized with a separate scale. // B is block window on shared memory // C is block distributed tensor -template +template struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { private: @@ -172,6 +157,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase using Base = BlockGemmAQuantBase; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -292,7 +278,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -302,7 +288,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { diff --git a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp index 7e28ea8fa9..077d0d8fe2 100644 --- a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp @@ -5,19 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" namespace ck_tile { -template +template struct BlockGemmBQuantBase { using BQDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; - static constexpr index_t UnaryOpSize = UnaryOpSize_; template CK_TILE_DEVICE static float cvt_scale_to_fp32(T scale) { @@ -42,24 +42,6 @@ struct BlockGemmBQuantBase } return scale_reg_f; } - - // can be inherited from A - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - const element_wise::PassThroughPack8 elementwise_op{}; - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } }; // A is block window on shared memory @@ -67,7 +49,9 @@ struct BlockGemmBQuantBase // Consecutive kQuantGroupSize elements of B are quantized with a separate scale. // B is block window on shared memory // C is block distributed tensor -template +template struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { private: @@ -170,6 +154,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase using Base = BlockGemmBQuantBase; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -291,7 +276,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -301,7 +286,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { diff --git a/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc b/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc index 035377734b..8f24c9bfe1 100644 --- a/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc +++ b/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc @@ -29,7 +29,8 @@ TYPED_TEST(TestCkTileBatchedGemm, Basic) {256, 256, 64, 8}, {256, 256, 64, 16}}; - if(ck_tile::get_device_name() != "gfx950") { + if(ck_tile::get_device_name() != "gfx950") + { gemmParams.emplace_back(256, 256, 128, 2); } diff --git a/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc b/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc index ab74e4e7b1..57feefceab 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc +++ b/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc @@ -2,6 +2,8 @@ // Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once +#include "ck_tile/host/permute_pk_int4.hpp" + template static constexpr inline auto is_row_major(Layout layout_) { @@ -91,61 +93,6 @@ void permute_tensor_b(Tensor& tensor) } } -template -void permute_vectors_i4x4_b(Tensor& tensor) -{ - const ck_tile::index_t K = tensor.get_length(0); - const ck_tile::index_t N = tensor.get_length(1); - // vector pk_i4x4 permute - for(int i = 0; i < N; i++) - { - for(int j = 0; j < K; j += 8) - { - int8_t input[8]; - - for(int k = 0; k < 4; k++) - { - int8_t i4x2 = tensor(j + k * 2, i).data; - input[k * 2 + 0] = (i4x2 >> 4) & 0xf; - input[k * 2 + 1] = (i4x2 >> 0) & 0xf; - } - - // permute 01234567->20643175 - { - int8_t hi = input[2]; - int8_t lo = input[0]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 0, i) = i4x2; - } - - { - int8_t hi = input[6]; - int8_t lo = input[4]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 2, i) = i4x2; - } - - { - int8_t hi = input[3]; - int8_t lo = input[1]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 4, i) = i4x2; - } - - { - int8_t hi = input[7]; - int8_t lo = input[5]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 6, i) = i4x2; - } - } - } -} - template ; -using WeightPreshuffle = - ck_tile::integral_constant; - -// Adding alias for the F8 parameters to facilitate skipping tests. -// This alias can be removed once test failures are fixed. -using F8Types = std::tuple; +using WeightPreshuffleV1 = + ck_tile::integral_constant; +using WeightPreshuffleV2 = + ck_tile::integral_constant; // clang-format off using KernelTypesWeightPreshuffle = ::testing::Types< - std::tuple< Row, Col, Row, F16, F16, F32, F16, Default, WeightPreshuffle>, - std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, Default, WeightPreshuffle> -#if !CK_TILE_USE_WMMA || CK_TILE_USE_OCP_FP8 - , F8Types + std::tuple< Row, Col, Row, F16, F16, F32, F16, Default, WeightPreshuffleV1>, + std::tuple< Row, Col, Row, F16, F16, F32, F16, Default, WeightPreshuffleV2>, + std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, Default, WeightPreshuffleV2>, + std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, Default, WeightPreshuffleV1> +#if !CK_TILE_USE_WMMA || CK_TILE_USE_OCP_FP8 + , + std::tuple< Row, Col, Row, F8, F8, F32, F16, Default, WeightPreshuffleV1>, + std::tuple< Row, Col, Row, F8, F8, F32, F16, Default, WeightPreshuffleV2>, + std::tuple< Row, Col, Row, F8, I4, F32, F16, Default, WeightPreshuffleV2>, + std::tuple< Row, Col, Row, F8, I4, F32, F16, Default, WeightPreshuffleV1> #endif >; diff --git a/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_ut_cases.inc b/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_ut_cases.inc index 389e0d53ea..bb56c63413 100644 --- a/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_ut_cases.inc +++ b/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_ut_cases.inc @@ -20,7 +20,7 @@ TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle) TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_128x128x128) { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v, F8>) { GTEST_SKIP() << "Skipping this test due to failures with F8"; } @@ -48,7 +48,7 @@ TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_128x128x4096) TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_128x2048x128) { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v, F8>) { GTEST_SKIP() << "Skipping this test due to failures with F8"; } @@ -77,7 +77,7 @@ TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_128x2048x4096) TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_1024x128x128) { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v, F8>) { GTEST_SKIP() << "Skipping this test due to failures with F8"; } @@ -106,7 +106,7 @@ TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_1024x128x4096) TYPED_TEST(TEST_SUITE_NAME, GemmPreshuffle_1024x2048x128) { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v, F8>) { GTEST_SKIP() << "Skipping this test due to failures with F8"; } diff --git a/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp index 42d0149498..62f819ac1e 100644 --- a/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp @@ -8,6 +8,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/host.hpp" #include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/permute_pk_int4.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" @@ -34,20 +35,31 @@ auto calculate_rtol_atol(const ck_tile::index_t K, enum struct GemmPipelineType { - WeightPreshuffle + WeightPreshuffleV1, + WeightPreshuffleV2 }; template struct GemmPipelineTypeSelector; template -struct GemmPipelineTypeSelector +struct GemmPipelineTypeSelector { using base_pipeline = ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV1; using pipeline = ck_tile::WeightPreshufflePipelineAGmemBGmemCRegV1; - static constexpr auto GetName() { return "GemmPipelineAgBgCrWeightPreshuffle"; } + static constexpr auto GetName() { return "GemmPipelineAgBgCrWeightPreshuffleV1"; } }; + +template +struct GemmPipelineTypeSelector +{ + using base_pipeline = ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2; + using pipeline = ck_tile::WeightPreshufflePipelineAGmemBGmemCRegV2; + + static constexpr auto GetName() { return "GemmPipelineAgBgCrWeightPreshuffleV2"; } +}; + template struct config { @@ -122,7 +134,8 @@ class TestCkTileGemmPipeline : public ::testing::Test constexpr bool kPadK = PadK; constexpr bool preshuffle = Preshuffle; - constexpr bool DoubleSmemBuffer = false; + constexpr bool DoubleSmemBuffer = + (PipelineType == GemmPipelineType::WeightPreshuffleV2) ? true : false; // TODO: For now - but this should also be a test parameter constexpr bool TransposeC = false; @@ -391,10 +404,19 @@ class TestCkTileGemmPipeline : public ::testing::Test ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - ck_tile::HostTensor b_shuffle_host = shuffle_b(b_k_n); - a_m_k_dev_buf.ToDevice(a_m_k.data()); - b_k_n_dev_buf.ToDevice(b_shuffle_host.data()); + ck_tile::HostTensor b_shuffle_host = shuffle_b(b_k_n); + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor b_shuffle_host_dev = b_shuffle_host; + ck_tile::permute_vectors_i4x4_b(b_shuffle_host_dev); + b_k_n_dev_buf.ToDevice(b_shuffle_host_dev.data()); + } + else + { + b_k_n_dev_buf.ToDevice(b_shuffle_host.data()); + } c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); From e469fee0460bb33cef2daa8ef9e05175b02195bc Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 18 Sep 2025 22:51:01 -0700 Subject: [PATCH 03/12] poc convert fnuz fp8 to non-native dtype similar to ocp (#2871) --- include/ck/utility/amd_ck_fp8.hpp | 30 +++++++++++++++++++++++++++-- include/ck/utility/data_type.hpp | 6 +++--- include/ck/utility/dtype_vector.hpp | 12 ++++++++++++ include/ck/utility/f8_utils.hpp | 29 ++++++++++++++-------------- include/ck/utility/type_convert.hpp | 12 ++++++------ test/data_type/test_bf8_fnuz.cpp | 20 ++++++++----------- test/data_type/test_fp8_fnuz.cpp | 20 ++++++++----------- 7 files changed, 80 insertions(+), 49 deletions(-) diff --git a/include/ck/utility/amd_ck_fp8.hpp b/include/ck/utility/amd_ck_fp8.hpp index 2edbb7c789..0b73f76155 100644 --- a/include/ck/utility/amd_ck_fp8.hpp +++ b/include/ck/utility/amd_ck_fp8.hpp @@ -33,8 +33,34 @@ namespace ck { -using f8_fnuz_t = _BitInt(8); -using bf8_fnuz_t = unsigned _BitInt(8); +struct f8_fnuz_t +{ + using data_type = unsigned char; + data_type m_data; + __host__ __device__ explicit constexpr f8_fnuz_t(data_type in_data) : m_data(in_data) {} + __host__ __device__ explicit constexpr f8_fnuz_t() = default; + __host__ __device__ bool constexpr operator==(f8_fnuz_t other) const + { + return m_data == other.m_data; + } + __host__ __device__ explicit constexpr operator data_type() const { return m_data; } +}; + +struct bf8_fnuz_t +{ + using data_type = unsigned char; + data_type m_data; + __host__ __device__ explicit constexpr bf8_fnuz_t(data_type in_data) : m_data(in_data) {} + __host__ __device__ explicit constexpr bf8_fnuz_t() = default; + __host__ __device__ bool constexpr operator==(bf8_fnuz_t other) const + { + return m_data == other.m_data; + } + __host__ __device__ explicit constexpr operator data_type() const { return m_data; } +}; + +static_assert(1 == sizeof(f8_fnuz_t)); +static_assert(1 == sizeof(bf8_fnuz_t)); typedef unsigned char fp8_storage_t; diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp index 48b352986e..984bb4d862 100644 --- a/include/ck/utility/data_type.hpp +++ b/include/ck/utility/data_type.hpp @@ -205,7 +205,7 @@ inline constexpr bool is_native_type() return is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || - is_same::value || is_same::value || is_same::value; + is_same_v || is_same_v || is_same::value; } // scalar_type @@ -300,14 +300,14 @@ struct scalar_type template <> struct scalar_type { - using type = f8_fnuz_t; + using type = f8_fnuz_t::data_type; static constexpr index_t vector_size = 1; }; template <> struct scalar_type { - using type = bf8_fnuz_t; + using type = bf8_fnuz_t::data_type; static constexpr index_t vector_size = 1; }; diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp index ae0edb35ee..27a7545a0e 100644 --- a/include/ck/utility/dtype_vector.hpp +++ b/include/ck/utility/dtype_vector.hpp @@ -1294,6 +1294,18 @@ struct nnvb_data_t_selector using type = bf8_ocp_t::data_type; }; +template <> +struct nnvb_data_t_selector +{ + using type = f8_fnuz_t::data_type; +}; + +template <> +struct nnvb_data_t_selector +{ + using type = bf8_fnuz_t::data_type; +}; + template <> struct nnvb_data_t_selector { diff --git a/include/ck/utility/f8_utils.hpp b/include/ck/utility/f8_utils.hpp index 799683ae65..748aa07f9e 100644 --- a/include/ck/utility/f8_utils.hpp +++ b/include/ck/utility/f8_utils.hpp @@ -39,7 +39,7 @@ __host__ __device__ Y run_cast_to_f8(X x, uint32_t rng) int exponent, bias; uint32_t head, mantissa, sign; // nan code is same for float and half - constexpr Y nan_code = 0x80; + constexpr uint8_t nan_code = 0x80; constexpr uint32_t nan_mask = NumericUtils::nan_mask; // convert to bitwise @@ -60,17 +60,17 @@ __host__ __device__ Y run_cast_to_f8(X x, uint32_t rng) if constexpr(negative_zero_nan) { if((x_bitwise & nan_mask) == nan_mask) - return nan_code; + return Y{nan_code}; } else { if((x_bitwise & nan_mask) == nan_mask) - return signed_inf + (mantissa != 0 ? 1 : 0); + return Y{static_cast(signed_inf + (mantissa != 0 ? 1 : 0))}; } // check if x is 0.0 if(x_bitwise == 0) - return 0; + return Y{0}; // First need to check if it is normal or denorm as there is a difference of implict 1 // Then need to adjust the exponent to align with the F8 exponent, in the meanwhile, shift @@ -178,9 +178,10 @@ In this case, the fp16 mantissa should be shift left by 1 */ // check if x is 0.0 or -0.0 if(out_exponent == 0 && mantissa == 0) - return negative_zero_nan ? 0 : (sign << (out_exp + out_mant)); + return Y{negative_zero_nan ? 0 : static_cast(sign << (out_exp + out_mant))}; mantissa &= (1 << out_mant) - 1; - return (sign << (out_exp + out_mant)) | (out_exponent << out_mant) | mantissa; + return Y{static_cast((sign << (out_exp + out_mant)) | (out_exponent << out_mant) | + mantissa)}; } template @@ -195,8 +196,8 @@ __host__ __device__ Y run_cast_from_f8(X x) constexpr int out_mant = NumericUtils::mant; // prepare the codes - constexpr X nan_code = 0x80; - using T_bitwise = typename NumericUtils::bitwise_type; + constexpr uint8_t nan_code = 0x80; + using T_bitwise = typename NumericUtils::bitwise_type; constexpr T_bitwise Inf_bitwise = NumericUtils::Inf; constexpr T_bitwise NegInf_bitwise = NumericUtils::NegInf; @@ -209,13 +210,13 @@ __host__ __device__ Y run_cast_from_f8(X x) constexpr Y Neg0 = bit_cast(Neg0_bitwise); // check if x is 0.0 - if(x == 0) + if(!static_cast(x)) return static_cast(0); // unpack the input - uint32_t sign = x >> (in_exp + in_mant); - uint32_t mantissa = x & ((1 << in_mant) - 1); - int exponent = (x & 0x7F) >> in_mant; + uint32_t sign = static_cast(x) >> (in_exp + in_mant); + uint32_t mantissa = static_cast(x) & ((1 << in_mant) - 1); + int exponent = (static_cast(x) & 0x7F) >> in_mant; constexpr int exp_low_cutoff = (1 << (out_exp - 1)) - (1 << (in_exp - 1)) + 1 - (negative_zero_nan ? 1 : 0); @@ -223,12 +224,12 @@ __host__ __device__ Y run_cast_from_f8(X x) if constexpr(negative_zero_nan) { - if(x == nan_code) + if(static_cast(x) == nan_code) return NaN; } else { - if(x == nan_code) + if(static_cast(x) == nan_code) return Neg0; if(exponent == ((1 << in_exp) - 1)) return (mantissa == 0) ? (sign ? NegInf : Inf) : NaN; diff --git a/include/ck/utility/type_convert.hpp b/include/ck/utility/type_convert.hpp index 290a6c8dd6..913557fc7a 100644 --- a/include/ck/utility/type_convert.hpp +++ b/include/ck/utility/type_convert.hpp @@ -351,7 +351,7 @@ inline __host__ __device__ f8_fnuz_t f8_convert_sr(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_fp8, -max_fp8); ival = __builtin_amdgcn_cvt_sr_fp8_f32(val.fval, rng, ival, 0); // 0 pos val.i32val = ival; - return val.i8val[0]; // little endian + return f8_t{val.i8val[0]}; // little endian #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -419,7 +419,7 @@ inline __host__ __device__ bf8_fnuz_t f8_convert_sr(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_bf8, -max_bf8); ival = __builtin_amdgcn_cvt_sr_bf8_f32(val.fval, rng, ival, 0); // 0 pos val.i32val = ival; - return val.i8val[0]; // little endian + return bf8_t{val.i8val[0]}; // little endian #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -655,7 +655,7 @@ inline __host__ __device__ f8_fnuz_t f8_convert_rne(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_fp8, -max_fp8); ival = __builtin_amdgcn_cvt_pk_fp8_f32(val.fval, val.fval, ival, false); // false -> WORD0 val.i32val = ival; - return val.i8val[0]; + return f8_t{val.i8val[0]}; #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -707,7 +707,7 @@ inline __host__ __device__ bf8_fnuz_t f8_convert_rne(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_bf8, -max_bf8); ival = __builtin_amdgcn_cvt_pk_bf8_f32(val.fval, val.fval, ival, false); // false -> WORD0 val.i32val = ival; - return val.i8val[0]; + return bf8_t{val.i8val[0]}; #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -924,7 +924,7 @@ inline __host__ __device__ float type_convert(f8_fnuz_t x) { #if defined(__gfx94__) float fval; - uint32_t i32val = static_cast(x); + uint32_t i32val = static_cast(static_cast(x)); fval = __builtin_amdgcn_cvt_f32_fp8(i32val, 0); // asm volatile("v_cvt_f32_fp8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); return fval; @@ -1430,7 +1430,7 @@ inline __host__ __device__ float type_convert(bf8_fnuz_t x) { #if defined(__gfx94__) float fval; - uint32_t i32val = static_cast(x); + uint32_t i32val = static_cast(static_cast(x)); fval = __builtin_amdgcn_cvt_f32_bf8(i32val, 0); // asm volatile("v_cvt_f32_bf8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); return fval; diff --git a/test/data_type/test_bf8_fnuz.cpp b/test/data_type/test_bf8_fnuz.cpp index 4ff796a614..f028c0da73 100644 --- a/test/data_type/test_bf8_fnuz.cpp +++ b/test/data_type/test_bf8_fnuz.cpp @@ -43,9 +43,8 @@ TEST(BF8FNUZ, ConvertFP32Nearest) type_convert(f8_convert_rne(std::numeric_limits::max())), abs_tol); // convert inf float to bf8_fnuz_t and check if it is qNan - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_rne(std::numeric_limits::infinity()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_rne(std::numeric_limits::infinity())); // positive norm float value to bf8 and back, check if holds float pos_float = 0.0000762939f; ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); @@ -80,9 +79,8 @@ TEST(BF8FNUZ, ConvertFP32Stochastic) type_convert(f8_convert_sr(std::numeric_limits::max())), abs_tol); // convert inf float to bf8_fnuz_t and check if it is qNan - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_sr(std::numeric_limits::infinity()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_sr(std::numeric_limits::infinity())); // positive norm float value to bf8 and back, check if holds float pos_float = 0.0000762939f; ASSERT_NEAR(pos_float, type_convert(f8_convert_sr(pos_float)), abs_tol); @@ -118,9 +116,8 @@ TEST(BF8FNUZ, ConvertFP16Nearest) type_convert(f8_convert_rne(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to bf8_fnuz_t and check if it is QuietNaN - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_rne(ck::NumericLimits::QuietNaN()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_rne(ck::NumericLimits::QuietNaN())); // positive norm fp16 value to bf8 and back, check if holds half_t pos_half = half_t{0.0000762939}; ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); @@ -155,9 +152,8 @@ TEST(BF8FNUZ, ConvertFP16Stochastic) type_convert(f8_convert_sr(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to bf8_fnuz_t and check if it is QuietNaN - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_sr(ck::NumericLimits::QuietNaN()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_sr(ck::NumericLimits::QuietNaN())); // positive norm fp16 value to bf8 and back, check if holds half_t pos_half = half_t{0.0000762939}; ASSERT_NEAR(pos_half, type_convert(f8_convert_sr(pos_half)), abs_tol); diff --git a/test/data_type/test_fp8_fnuz.cpp b/test/data_type/test_fp8_fnuz.cpp index c2ec6dad94..0cf775f947 100644 --- a/test/data_type/test_fp8_fnuz.cpp +++ b/test/data_type/test_fp8_fnuz.cpp @@ -48,9 +48,8 @@ TEST(FP8FNUZ, ConvertFP32Nearest) type_convert(f8_convert_rne(std::numeric_limits::max())), abs_tol); // convert inf float to f8_fnuz_t and check if it is qNan - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_rne(std::numeric_limits::infinity()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_rne(std::numeric_limits::infinity())); // positive norm float value to fp8 and back, check if holds float pos_float = 0.017578125f; ASSERT_NEAR(pos_float, type_convert(f8_convert_rne(pos_float)), abs_tol); @@ -85,9 +84,8 @@ TEST(FP8FNUZ, ConvertFP32Stochastic) type_convert(f8_convert_sr(std::numeric_limits::max())), abs_tol); // convert inf float to f8_fnuz_t and check if it is qNan - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_sr(std::numeric_limits::infinity()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_sr(std::numeric_limits::infinity())); // positive norm float value to fp8 and back, check if holds float pos_float = 0.017578125f; ASSERT_NEAR(pos_float, type_convert(f8_convert_sr(pos_float)), abs_tol); @@ -122,9 +120,8 @@ TEST(FP8FNUZ, ConvertFP16Nearest) type_convert(f8_convert_rne(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to f8_fnuz_t and check if it is QuietNaN - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_rne(ck::NumericLimits::QuietNaN()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_rne(ck::NumericLimits::QuietNaN())); // positive norm fp16 value to fp8 and back, check if holds half_t pos_half = half_t{0.017578125}; ASSERT_NEAR(pos_half, type_convert(f8_convert_rne(pos_half)), abs_tol); @@ -159,9 +156,8 @@ TEST(FP8FNUZ, ConvertFP16Stochastic) type_convert(f8_convert_sr(ck::NumericLimits::Max())), abs_tol); // convert QuietNaN fp16 to f8_fnuz_t and check if it is QuietNaN - ASSERT_NEAR(ck::NumericLimits::QuietNaN(), - f8_convert_sr(ck::NumericLimits::QuietNaN()), - abs_tol); + ASSERT_EQ(ck::NumericLimits::QuietNaN(), + f8_convert_sr(ck::NumericLimits::QuietNaN())); // positive norm fp16 value to fp8 and back, check if holds half_t pos_half = half_t{0.017578125}; ASSERT_NEAR(pos_half, type_convert(f8_convert_sr(pos_half)), abs_tol); From dd249f1cd6c516f7a1d45663f7f26eb4a4c086ca Mon Sep 17 00:00:00 2001 From: ltqin Date: Fri, 19 Sep 2025 14:26:43 +0800 Subject: [PATCH 04/12] Add input fp8 and output bf16 attention (#2726) * change host using fp16 to check * fp8 to fp8 compare * rewrite input parameters * add not squant * remove some output code * for scale = 1 * format * saturates only for fp8 * add fp8bf16 data type * add fp8bf16 data type * fix test fp8 code * add run_fp8bf16_tests * change fmha fwd example parameter(adding fp8bf16) * Support fp8bf16 for Aiter * Support aiter fp8bf16 in c++ * fix comment about fp8 in readme.md * add fp8fp32 * add fp8fp32 test * remove range_q etc. * format * fix test parameters about squant and fmha example input fp8bf16 fp8fp32 data type * add fp8bf16 to data_type function * change colmajor to rowmajor in test_ck_tile_fmha_fwd_fp8 * format * reset atol for fp8 * fix bug for atol --------- Co-authored-by: rocking Co-authored-by: asleepzzz --- example/ck_tile/01_fmha/README.md | 2 +- .../ck_tile/01_fmha/codegen/cpp_symbol_map.py | 3 +- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 48 +++-- .../01_fmha/codegen/ops/fmha_fwd_splitkv.py | 1 - .../codegen/ops/fmha_pagedkv_prefill.py | 10 +- example/ck_tile/01_fmha/example_fmha_fwd.cpp | 30 ++-- example/ck_tile/01_fmha/fmha_fwd.hpp | 36 ++++ example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 170 +++++++++++------- .../ck_tile/01_fmha/script/smoke_test_fwd.sh | 29 ++- .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 52 +++--- test/ck_tile/fmha/test_fmha_fwd.inc | 21 +-- test/ck_tile/fmha/test_fmha_fwd_fp8.cpp | 13 +- 12 files changed, 262 insertions(+), 153 deletions(-) diff --git a/example/ck_tile/01_fmha/README.md b/example/ck_tile/01_fmha/README.md index cb6cd44f64..7f55d7412f 100644 --- a/example/ck_tile/01_fmha/README.md +++ b/example/ck_tile/01_fmha/README.md @@ -131,4 +131,4 @@ TBD ## FP8 experimental support As described in [this blog](https://blog.hippoml.com/8bit-hippoattention-up-to-3x-faster-compared-to-flashattentionv2-8f9def90b482), we have an experimental support for fp8 fmha kernels, you can evaluate the performance by setting the arg `-prec=fp8` to the `tile_example_fmha_fwd`, on a gfx942 machine and ROCm 6.0+. -Currently we only support `-vlayout=c`( `hdim*seqlen` for V matrix) and `-squant=1`(static quantization) with `hdim=128` for fp8 now. Full feature support will come later. +Currently we only support `-vlayout=r`( `seqlen*hdim` for V matrix) for fp8 and fp8bf16 now. Full feature support will come later. diff --git a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py index 42a9d5148a..802c9e51d7 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -7,7 +7,8 @@ FWD_DTYPE_MAP = { "bf16" : "FmhaFwdBf16", "fp8" : "FmhaFwdFp8", "fp8fp16": "FmhaFwdFp8Fp16", - "fp8bf16": "FmhaFwdFp8Bf16" + "fp8bf16": "FmhaFwdFp8Bf16", + "fp8fp32": "FmhaFwdFp8Fp32" } BWD_DTYPE_MAP = { diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index d9452206e7..cfb96b7d53 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -163,7 +163,7 @@ float fmha_fwd(fmha_fwd_traits t, fmha_fwd_args a, const ck_tile::stream_config& [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) {{ return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); }}; - + const bool has_load_tr = ck_tile::is_load_tr_supported(); {F_dispatch} @@ -248,11 +248,11 @@ class FmhaFwdApiTrait: if self.spad == 't' : return f'true /*a.seqlen_q % {self.bm0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.seqlen_q % {self.bm0} == 0' else: assert False - + @property def seqtune(self) -> str: if self.bm0 == 128: return 'true/*fall back to largest tile*/' # group mode only generate spad/skpad == true - else: + else: return f'a.seqlen_q <= {self.bm0}' @property @@ -351,7 +351,7 @@ class FmhaFwdPipeline: if self.F_squant == 't' : n += '_squant' else: n += '_nsquant' - + if self.F_trload == 't' : n += '_trload' else: n += '_ntrload' @@ -378,7 +378,7 @@ class FmhaFwdApiPool: "t": "has_load_tr", "f": "true" } - + per_tr_load =str() for tr_load in ["t", "f"]: per_dtypes=str() @@ -550,12 +550,16 @@ class KernelComponentFactory: (192,192) : [FmhaFwdTileSize(128, 128, 32, 192, 32, 192, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, 1)], (256,256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], } - elif dtype == 'fp8' or dtype == 'bf8': + elif dtype == 'fp8' or dtype == 'fp8bf16': return { (64,64 ) : [FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (128,128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (256,256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], } + elif dtype == 'fp8fp32': + return { + (128,128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], + } else: return None @@ -567,9 +571,9 @@ class KernelComponentFactory: # TODO: the order of List matters! the later in this list will be also be checked later # TODO: currently for qr pipeline, let 't' padding to appear later!! # TODO: how to design this more generic? - squant = 't' if dtype == 'fp8' else 'f' pipelines = [] if dtype in ['fp16', 'bf16']: + squant = 'f' for logits, mask, bias, lse, dropout, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]): if hdim == 256 and hdim_v == 256: pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, skip, 'f')) @@ -589,11 +593,12 @@ class KernelComponentFactory: pipelines.append(FmhaFwdPipeline('qr_async_trload', 'row', 'f', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 't')) if receipt == 1 and bias != "bias": pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 'f')) # TODO: cover arbitraty hdim - elif dtype in ['fp8', 'bf8']: + elif dtype in ['fp8', 'fp8bf16', 'fp8fp32']: # no need lse/dropout kernels - for logits, mask, bias in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): - pipelines.append(FmhaFwdPipeline('qr', 'col', 'f', 'f', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) - elif dtype in ['fp8fp16', 'fp8bf16']: + for logits, squant, mask, bias in itertools.product(["f"], ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): + pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) + pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) + elif dtype in ['fp8fp16', 'bf8']: # TODO None else: @@ -674,25 +679,34 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl continue # Aiter(mha_fwd) integration elif receipt == 100: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= mode == 'batch' cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 if not cond: continue # Aiter(mha_varlen_fwd) integration elif receipt == 200: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= mode == 'group' cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 if not cond: continue # aiter::mha_fwd C++ api integration elif receipt == 600: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 + if not cond: + continue + elif receipt == 888: + cond = dtype in ['fp8', 'fp8bf16', 'fp8fp32'] + cond &= pipeline.F_vlayout == 'row' + cond &= hdim == 128 if not cond: continue diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index 3b48b3d005..cee1505486 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -645,7 +645,6 @@ def get_fmha_fwd_tile_dict_from_dtype(dtype : str) -> Optional[dict]: return { '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1), '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), - '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), } else: return None diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py index 7b93e9654c..df6b422981 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py @@ -465,14 +465,14 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl squant = 't' if dtype == 'fp8' else 'f' pipelines = [] if dtype in ['fp16', 'bf16']: - for logits, mask, bias, pagedkv, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"]): - pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'col', 't', 'f', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) - pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'col', 't', 't', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) + for logits, mask, bias, pagedkv, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t"], ["f"]): pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 'f', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 't', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) elif dtype in ['fp8', 'bf8']: - # TODO - None + # no need lse/dropout kernels + for logits, mask, bias in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): + pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 'f', 'f', 'f', 'f', logits, bias, 'f', 't', squant, mask, 'f')) + pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 't', 'f', 'f', logits, bias, 'f', 't', squant, mask, 'f')) elif dtype in ['fp8fp16', 'fp8bf16']: # TODO None diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp index c3bbb7a558..91cb9f55be 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp @@ -44,21 +44,15 @@ auto create_args(int argc, char* argv[]) .insert("scale_s", "0", "scale factor of S. 0 means equal to 1/sqrt(hdim).\n" - "note when squant=1, this value will be modified by range_q/k") + "note when squant=1, this value will be modified") .insert("logits_soft_cap", "0", "attention logits soft capping value.") - .insert("range_q", "16", "per-tensor quantization range of q. used if squant=1.") - .insert("range_k", "16", "per-tensor quantization range of k. used if squant=1.") - .insert("range_v", "16", "per-tensor quantization range of v. used if squant=1.") - .insert("range_p", "1", "per-tensor quantization range of p [e^(s-m)]. used if squant=1.") - .insert("range_o", "16", "per-tensor quantization range of o (p*v). used if squant=1.") .insert("squant", "auto", "if using static quantization fusion or not. auto: fp8 will default use squant, " "other will not\n" "0: no static quant(not implemented) 1: apply scale_p and scale_o with respect to " "P and O.\n" - "calculate scale_s, scale_p, scale_o according to range_q, range_k, range_v, " - "range_p, range_o") + "calculate scale_s, scale_p, scale_o auto") .insert("iperm", "1", "permute input\n" @@ -89,7 +83,7 @@ auto create_args(int argc, char* argv[]) "uf", "init method:\n ui or 0 - uniform random int\n ni - normalized random int" "\n uf or 1 - uniform random float\n nf - normalized random float" - "\n tf or 2 - trig float\n uf:q or ufq or 3 - fp8 quantization") + "\n tf or 2 - trig float\n") .insert("seed", "11939", "random seed used for initializing input tensors. 0 for " @@ -148,11 +142,6 @@ auto run(const ck_tile::ArgParser& arg_parser) uint64_t drop_offset = arg_parser.get_uint64("drop_offset"); bool drop_prefs = arg_parser.get_bool("drop_prefs"); std::string mask_str = arg_parser.get_str("mask"); - float range_q = arg_parser.get_float("range_q"); - float range_k = arg_parser.get_float("range_k"); - float range_v = arg_parser.get_float("range_v"); - float range_p = arg_parser.get_float("range_p"); - float range_o = arg_parser.get_float("range_o"); bool is_rotary_interleaved = arg_parser.get_bool("rotary_interleaved"); ck_tile::index_t num_splits = arg_parser.get_int("num_splits"); std::string init_method = arg_parser.get_str("init"); @@ -201,11 +190,6 @@ auto run(const ck_tile::ArgParser& arg_parser) drop_offset, drop_prefs, mask_str, - range_q, - range_k, - range_v, - range_p, - range_o, squant, is_rotary_interleaved, num_splits, @@ -237,6 +221,14 @@ int main(int argc, char* argv[]) { return run(arg_parser) == fwd_result::success ? 0 : -2; } + else if(data_type == "fp8bf16") + { + return run(arg_parser) == fwd_result::success ? 0 : -2; + } + else if(data_type == "fp8fp32") + { + return run(arg_parser) == fwd_result::success ? 0 : -2; + } std::cerr << "Unsupported precision: " << data_type << std::endl; return -1; } diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index df1e9e5699..c41e48e6aa 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -41,6 +41,10 @@ struct FmhaFwdFp8Bf16 { }; +struct FmhaFwdFp8Fp32 +{ +}; + template struct FmhaFwdTypeConfig; @@ -108,6 +112,38 @@ struct FmhaFwdTypeConfig using ODataType = ck_tile::bf8_t; }; +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::bf16_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; +}; + struct FmhaMasks { using NoMask = ck_tile::GenericAttentionMask; diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 397245ab32..43f484fe14 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -50,20 +50,30 @@ auto get_elimit(std::string /*init_method*/) } template <> -auto get_elimit(std::string init_method) +auto get_elimit(std::string /*init_method*/) { - if(init_method == "ui" || init_method == "ni") - { - unsigned max_rounding_point_distance = 0; - double atol = 2e-3; - return ck_tile::make_tuple(max_rounding_point_distance, atol); - } - else - { - unsigned max_rounding_point_distance = 1; - double atol = 0.0625; - return ck_tile::make_tuple(max_rounding_point_distance, atol); - } + using TypeConfig = FmhaFwdTypeConfig; + using ODataType = typename TypeConfig::ODataType; + float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + double rtol = 0; + double atol = 16 * (o_dtype_max > 240 ? 2 : 1); + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-2; + double atol = 1.8e-1; + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-2; + double atol = 1.8e-1; + return ck_tile::make_tuple(rtol, atol); } int num_splits_heuristic(int batch_nhead_mblocks, int num_SMs, int max_splits) @@ -157,11 +167,6 @@ fwd_result fmha_fwd_run(mode_enum mode, uint64_t drop_offset, bool drop_prefs, std::string mask_str, - float range_q, - float range_k, - float range_v, - float range_p, - float range_o, bool squant, bool is_rotary_interleaved, ck_tile::index_t num_splits, @@ -180,6 +185,10 @@ fwd_result fmha_fwd_run(mode_enum mode, return "fp8"; else if constexpr(std::is_same_v) return "bf8"; + else if constexpr(std::is_same_v) + return "fp8bf16"; + else if constexpr(std::is_same_v) + return "fp8fp32"; else static_assert(false); }(); @@ -367,22 +376,6 @@ fwd_result fmha_fwd_run(mode_enum mode, using OaccDataType = typename TypeConfig::OaccDataType; using ODataType = typename TypeConfig::ODataType; - float q_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float k_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float v_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float p_dtype_max = v_dtype_max; // assume p and v is the same type - float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - - float scale_p = 1.f; - float scale_o = 1.f; - - if(squant) - { - scale_s = scale_s * (range_q / q_dtype_max) * (range_k / k_dtype_max); - scale_p = p_dtype_max / range_p; - scale_o = (o_dtype_max / range_o) * (range_p / p_dtype_max) * (range_v / v_dtype_max); - } - // accumulation numbers for performance evaluation std::size_t flop = 0, num_byte = 0; auto max_seqlen_q = @@ -528,7 +521,7 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::HostTensor cache_batch_idx_host(use_cache_batch_idx ? std::array{batch} : std::array{1}); - + float max_o = 5.0; if(init_method == "ui" || init_method == "0") { ck_tile::FillUniformDistributionIntegerValue{-3.f, 3.f, next_seed()}(q_host); @@ -576,32 +569,6 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::FillTrigValue{}(vnew_host); ck_tile::FillTrigValue{}(bias_host); } - else if(init_method == "ufq" || init_method == "uf:q" || init_method == "3") - { - // suitable for fp8 quantization - if(!squant) - { - std::cerr << "init method " << init_method << " can not be used without quantization" - << std::endl; - return fwd_result::invalid_args; - } - ck_tile::FillUniformDistribution{0.f, q_dtype_max, next_seed()}(q_host); - ck_tile::FillUniformDistribution{0.f, k_dtype_max, next_seed()}(k_host); - ck_tile::FillUniformDistribution{0.f, k_dtype_max, next_seed()}(knew_host); - ck_tile::FillUniformDistribution{0.f, v_dtype_max, next_seed()}(v_host); - ck_tile::FillUniformDistribution{0.f, v_dtype_max, next_seed()}(vnew_host); - - // bias_fp8 = qscale_bias * bias_fp32 - float qscale_bias = (q_dtype_max / range_q) * (k_dtype_max / range_k); - // Assume bias is in [0.f, 1.f] in original fp32 - ck_tile::FillUniformDistribution{0.f, qscale_bias, next_seed()}(bias_host); - } - else - { - std::cerr << "Unknown value for init argument: " << init_method << std::endl; - return fwd_result::invalid_args; - } - if(bias.type == bias_enum::alibi) { auto slopes = ck_tile::get_alibi_slopes(nhead); @@ -625,8 +592,8 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem q_buf(q_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem k_buf(k_host.get_element_space_size_in_bytes()); - ck_tile::DeviceMem knew_buf(knew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem v_buf(v_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem knew_buf(knew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem vnew_buf(vnew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem bias_buf(bias_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem lse_acc_buf(lse_acc_host.get_element_space_size_in_bytes()); @@ -650,10 +617,79 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem block_table_buf(block_table_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem cache_batch_idx_buf(cache_batch_idx_host.get_element_space_size_in_bytes()); + float scale_p = 1.f; + float scale_o = 1.f; + if(squant) + { + float q_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float p_dtype_max = v_dtype_max; // assume p and v is the same type + // Q tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + q_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + + float scale = q_dtype_max / max_value; + + q_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + scale_s = scale_s / scale; + } + + // K tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + k_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + float scale = k_dtype_max / max_value; + k_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + scale_s = scale_s / scale; + } + + // V tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + v_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + + float scale = k_dtype_max / max_value; + v_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + + scale_o = (1.0 / p_dtype_max) / scale; + } + + scale_p = p_dtype_max; + + if constexpr(std::is_same_v) + { + float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + scale_o = scale_o * o_dtype_max / max_o; + } + } + q_buf.ToDevice(q_host.data()); k_buf.ToDevice(k_host.data()); - knew_buf.ToDevice(knew_host.data()); v_buf.ToDevice(v_host.data()); + knew_buf.ToDevice(knew_host.data()); vnew_buf.ToDevice(vnew_host.data()); bias_buf.ToDevice(bias_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); @@ -1103,7 +1139,9 @@ fwd_result fmha_fwd_run(mode_enum mode, lse_buf.FromDevice(lse_host.data()); randval_buf.FromDevice(randval_host.data()); - constexpr bool supports_squant = std::is_same_v; + constexpr bool supports_squant = std::is_same_v || + std::is_same_v || + std::is_same_v; auto p_compute_element_func = [&]() { if constexpr(supports_squant) @@ -1113,9 +1151,11 @@ fwd_result fmha_fwd_run(mode_enum mode, }(); auto oacc_element_func = [&]() { - if constexpr(supports_squant) + if constexpr(std::is_same_v && supports_squant) return ck_tile::composes(ck_tile::saturates{}, ck_tile::scales{scale_o}); + else if constexpr(supports_squant) + return ck_tile::scales{scale_o}; else return ck_tile::identity{}; }(); diff --git a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh index c087a1fb3e..afd0c728c6 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh @@ -94,7 +94,30 @@ run_fp8_tests() { for b in 1 2 ; do for hdim in 64 128 256 ; do - run_exe -prec=fp8 -init=3 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=c -squant=1 -kname=$KNAME $COMMON_ARGS + $EXE -prec=fp8 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + + done ; done ; done ; done +} + +run_fp8bf16_tests() { + for perm in 0 1 ; do + for bias in "n" "e" "a" ; do + for b in 1 2 ; do + for hdim in 64 128 256 ; do + + $EXE -prec=fp8bf16 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + + done ; done ; done ; done +} + +run_fp8fp32_tests() { + for perm in 0 1 ; do + for bias in "n" "e" "a" ; do + for b in 1 2 ; do + for hdim in 64 128 256 ; do + + $EXE -prec=fp8fp32 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + done ; done ; done ; done } @@ -117,7 +140,9 @@ run_fp16_appendkv_tests() { set -x run_fp16_bf16_tests -# run_fp8_tests +run_fp8_tests +run_fp8bf16_tests +run_fp8fp32_tests if [ $TEST_APPENDKV -eq 1 ] ; then run_fp16_appendkv_tests diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 6405ca50df..58fdad149a 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -1446,29 +1446,35 @@ struct FmhaFwdKernel auto o_acc_tile = [&]() { if constexpr(kDoFp8StaticQuant) { - return FmhaPipeline{}( - q_dram_window, - identity{}, // q_element_func - k_dram_window, - identity{}, // k_element_func - v_dram_window, - identity{}, // v_element_func - bias_dram_window, - identity{}, // bias_element_func - randval_dram_window, - lse_dram_window, - identity{}, // lse_element_func - identity{}, // s_acc_element_func - scales{kargs.scale_p}, // p_compute_element_func - composes(saturates{}, scales{kargs.scale_o}), // o_acc_element_func - mask, - position_encoding, - kargs.scale_s, - variant, - variant_params, - block_indices, - smem_ptr, - dropout); + auto o_acc_element_func = [&]() { + if constexpr(std::is_same_v) + return ck_tile::composes(ck_tile::saturates{}, + ck_tile::scales{kargs.scale_o}); + else + return ck_tile::scales{kargs.scale_o}; + }(); + return FmhaPipeline{}(q_dram_window, + identity{}, // q_element_func + k_dram_window, + identity{}, // k_element_func + v_dram_window, + identity{}, // v_element_func + bias_dram_window, + identity{}, // bias_element_func + randval_dram_window, + lse_dram_window, + identity{}, // lse_element_func + identity{}, // s_acc_element_func + scales{kargs.scale_p}, // p_compute_element_func + o_acc_element_func, // o_acc_element_func + mask, + position_encoding, + kargs.scale_s, + variant, + variant_params, + block_indices, + smem_ptr, + dropout); } else { diff --git a/test/ck_tile/fmha/test_fmha_fwd.inc b/test/ck_tile/fmha/test_fmha_fwd.inc index f02ef1e55e..08abd3358d 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.inc +++ b/test/ck_tile/fmha/test_fmha_fwd.inc @@ -32,9 +32,6 @@ const ck_tile::stream_config stream_config{ 1, // rotating_count_ }; -// range_q, range_k, range_v, range_p, range_o, squant -#define QUANT_ARGS 1, 1, 1, 1, 1, squant - #define COMMON_ARGS \ init_method, static_cast(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))), 1, \ stream_config @@ -117,7 +114,7 @@ TEST_P(AllLong, Test) 1024, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -179,7 +176,7 @@ TEST_P(HDimPadding, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -236,7 +233,7 @@ TEST_P(ElementwiseBias, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -292,7 +289,7 @@ TEST_P(Alibi, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -350,7 +347,7 @@ TEST_P(Dropout, Test) drop_offset, // drop_offset drop_prefs, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -410,7 +407,7 @@ TEST_P(PagedKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -476,7 +473,7 @@ TEST_P(SplitKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved num_splits, // num_splits COMMON_ARGS); @@ -548,7 +545,7 @@ TEST_P(AppendKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, false, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -618,7 +615,7 @@ TEST_P(AppendKVRoPE, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, is_rotary_interleaved, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); diff --git a/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp b/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp index 46ed8f4125..b99c304d1f 100644 --- a/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp +++ b/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp @@ -17,22 +17,21 @@ using DataTypeConfig = FmhaFwdFp8; // instances are added), however the corresponding tests are not disabled (they will be skipped) // in case such instances will be added in the future. -const auto HDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto HDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); -const auto SplitKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto SplitKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); -const auto AppendKVHDimValues = - Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto AppendKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); // There are no fp8 instances with seqlen padding (mode_enum::group requires it) const auto ModeValues = Values(mode_enum::batch); const auto IsVRowmajorValues = Values(false); -const bool squant = true; -const std::string init_method = "ufq"; +const auto squant = true; +const std::string init_method = "uf"; const bool def_lse = false; -const bool def_is_v_rowmajor = false; +const bool def_is_v_rowmajor = true; int adjust_seqlen(int seqlen) { From 2aec38f9ec67bfbdccbdb3a5c25913e5a9ba6136 Mon Sep 17 00:00:00 2001 From: Anton Gorenko Date: Fri, 19 Sep 2025 12:34:45 +0600 Subject: [PATCH 05/12] [CK_TILE] FMHA Fix synchronization issues in BWD pipelines (#2876) * Run ctest with --output-on-failure * Fix synchronization issues in bwd pipelines The bwd kernel reuses the same area of LDS for ds (SGrad), bias and dbias (BiasGrad). This means that there must be block_sync_lds between loading one tensor and storing another to the same area. Heavy instructions like MFMA/WMMA and global loads are executed between reuses of the same memory so in MOST cases loading is finished by all warps before storing is started. However, sometimes warps progress at different speeds. Running the tests multiple times and, preferably, with multiple processes on the same GPU helps to trigger this issue: bin/test_ck_tile_fmha_bwd_bf16 --gtest_repeat=-1 --gtest_shuffle --gtest_throw_on_failure --- ...fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp | 16 +++++++++++----- ...ha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp | 6 ++++++ ...a_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp | 6 ++++++ .../block_fmha_bwd_pipeline_default_policy.hpp | 2 +- script/launch_tests.sh | 4 +--- 5 files changed, 25 insertions(+), 9 deletions(-) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp index b883aad155..c402eaeac4 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp @@ -559,6 +559,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP auto shuffled_bias_tile = make_static_distributed_tensor( Policy::template MakeShuffledBiasTileDistribution()); shuffle_tile(shuffled_bias_tile, bias_tile); + // SGrad and Bias use the same address in LDS, finish loading ds on the previous + // iteration to reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, shuffled_bias_tile); block_sync_lds(); auto bias_s_tile = load_tile(bias_s_lds_read_window); @@ -814,6 +817,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP auto shuffled_bias_tile = make_static_distributed_tensor( Policy::template MakeShuffledBiasTileDistribution()); shuffle_tile(shuffled_bias_tile, bias_tile); + // SGrad and Bias use the same address in LDS, finish loading ds in the hot loop to + // reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, shuffled_bias_tile); block_sync_lds(); auto bias_s_tile = load_tile(bias_s_lds_read_window); @@ -956,6 +962,8 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP return cast_tile(ds); } }(); + // Finish loading bias_s to reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, dbias); block_sync_lds(); auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); @@ -975,11 +983,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); - if constexpr(kHasBiasGrad) - { - // SGrad and BiasGrad use the same address in LDS. - block_sync_lds(); - } + // SGrad and Bias/BiasGrad use the same address in LDS, finish loading bias/dbias or, when + // bias is not used, loading ds in the hot loop to reuse LDS. + block_sync_lds(); store_tile(ds_lds_window, ds_gemm); block_sync_lds(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp index 81950bd30a..41cb4fc306 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp @@ -698,6 +698,12 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadKRKTRVR dst_reg_tensor.get_thread_buffer() = ds_gemm.get_thread_buffer(); gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); + if constexpr(kHasBiasGrad) + { + // SGrad and BiasGrad use the same address in LDS, finish loading dbias to reuse + // LDS. + block_sync_lds(); + } store_tile(ds_lds_window, ds_gemm); } s_waitcnt(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp index 16d9f695df..8c8d2af486 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp @@ -656,6 +656,12 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR dst_reg_tensor.get_thread_buffer() = ds_gemm.get_thread_buffer(); dk_acc = gemm_3(dst_reg_tensor, qt_reg_tensor); + if constexpr(kHasBiasGrad) + { + // SGrad and BiasGrad use the same address in LDS, finish loading dbias to reuse + // LDS. + block_sync_lds(); + } store_tile(ds_lds_window, ds_gemm); } __builtin_amdgcn_s_waitcnt(3952); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp index 68ead7c765..ad9e2959f5 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp @@ -1941,7 +1941,7 @@ struct BlockFmhaBwdPipelineDefaultPolicy constexpr index_t smem_size_stage0_0 = smem_size_k + smem_size_kt; constexpr index_t smem_size_stage0_1 = smem_size_v; - constexpr index_t smem_size_stage1 = smem_size_qt + smem_size_q + +smem_size_dot + + constexpr index_t smem_size_stage1 = smem_size_qt + smem_size_q + smem_size_dot + smem_size_do + smem_size_lse + smem_size_d + max(smem_size_bias, smem_size_ds); diff --git a/script/launch_tests.sh b/script/launch_tests.sh index 5e71e25478..17a99e62a3 100755 --- a/script/launch_tests.sh +++ b/script/launch_tests.sh @@ -49,7 +49,7 @@ with open('$TEST_FILE', 'r') as f: if tests: # Extract just the filename after the last '/' clean_tests = [os.path.basename(test) for test in tests] - print('ctest -R \"' + '|'.join(clean_tests) + '\"') + print('ctest --output-on-failure -R \"' + '|'.join(clean_tests) + '\"') else: print('# No tests to run') ") @@ -57,5 +57,3 @@ with open('$TEST_FILE', 'r') as f: echo "$command" eval "$command" - - From 86dd59cd01e41a4190bf2405a0fb0e89d9498b4c Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Fri, 19 Sep 2025 17:36:49 +0800 Subject: [PATCH 06/12] =?UTF-8?q?[CK=5FTILE]=20Add=20sequence=20padding=20?= =?UTF-8?q?and=20variable=20length=20support=20in=20fmha=20(a=E2=80=A6=20(?= =?UTF-8?q?#2851)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [CK_TILE] Add sequence padding and variable length support in fmha (and v3) - Group Mode Padding: Introduces the `-s_qpad` argument to support physically padded layouts. Kernels now use padded start pointers (`seqstart_padded_*_ptr`) for memory addressing. - Batch Mode Variable Length: Adds `-q_eff_lens` and `-kv_eff_lens` arguments for efficient processing of variable-length sequences by passing cumulative effective lengths (`cu_seqlen_*_ptr`) to the kernel. - FMHA examples: Support padding and variable length both in group and batch mode. Dispatcher is updated as well (dispatch to kPadSeqLenK enabled pipeline). - New padding test cases: Add padding test cases to `smoke_test_fwd.sh`, and add benchmarks to `benchmark_fwd.sh` and `benchmark_fwd_v3.sh` as well. These test cases and benchmarks that specifically validate/benchmark the new padding and variable-length functionalities in both group and batch modes. * [CK_TILE] Fix build error in fmha unit tests --------- Co-authored-by: Po Yen Chen Co-authored-by: Yi DING --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 6 +- example/ck_tile/01_fmha/example_fmha_fwd.cpp | 20 +- .../ck_tile/01_fmha/example_fmha_fwd_v3.cpp | 148 ++++++++- example/ck_tile/01_fmha/fmha_fwd.hpp | 17 +- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 127 +++++++- example/ck_tile/01_fmha/fmha_fwd_v3.hpp | 5 + example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp | 4 +- .../ck_tile/01_fmha/script/benchmark_fwd.sh | 33 ++ .../01_fmha/script/benchmark_fwd_v3.sh | 17 ++ .../ck_tile/01_fmha/script/smoke_test_fwd.sh | 109 +++++++ .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 285 ++++++++++++++++-- .../ops/fmha/kernel/fmha_fwd_v3_kernel.hpp | 180 ++++++++++- test/ck_tile/fmha/test_fmha_fwd.inc | 141 +++++++++ 13 files changed, 1032 insertions(+), 60 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index cfb96b7d53..da0c9ca931 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -259,11 +259,11 @@ class FmhaFwdApiTrait: def skcheck(self) -> str: if self.mode == 'group': return 'true/*group mode skpad always true*/' # group mode only generate spad/skpad == true if self.pipeline_tag == 'qr_async': - if self.skpad == 't' : return f'a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0' - else : return f'a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0' + if self.skpad == 't' : return f'(a.cu_seqlen_kv_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0)' + else : return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)' elif self.pipeline_tag in ['qr', 'qs']: if self.skpad == 't' : return f'true /*a.seqlen_k % {self.bn0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) - else : return f'a.seqlen_k % {self.bn0} == 0' + else : return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)' elif self.pipeline_tag == 'qr_async_trload': if self.skpad == 't' : return 'true' else: return 'true' diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp index 91cb9f55be..79fda6d564 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp @@ -33,6 +33,10 @@ auto create_args(int argc, char* argv[]) "0", "seqlen_k for new key/value, 0 means not to use this at all; " "-1 to choose s_knew in [1, s] randomly.") + .insert("s_qpad", + "-1", + "seqlen_q stride between 2 batches (group-mode optional).\n" + "Provide positive strides per-batch to simulate physical padding on Q.") .insert("s_kpad", "-1", "seqlen_k stride between 2 batches, currently used in group-mode only\n" @@ -107,7 +111,15 @@ auto create_args(int argc, char* argv[]) .insert("warmup", "5", "number of iterations before benchmark the kernel") .insert("repeat", "20", "number of iterations to benchmark the kernel") .insert("json", "0", "0: No Json, 1: Dump Results in Json format") - .insert("jsonfile", "fmha_fwd.json", "json file name to dump results"); + .insert("jsonfile", "fmha_fwd.json", "json file name to dump results") + .insert("q_eff_lens", + "", + "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n" + "Comma-separated list of length 'b'. If empty, no override.") + .insert("kv_eff_lens", + "", + "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n" + "Comma-separated list of length 'b'. If empty, no override."); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); @@ -127,6 +139,9 @@ auto run(const ck_tile::ArgParser& arg_parser) ck_tile::index_t hdim_v = arg_parser.get_int("d_v"); ck_tile::index_t seqlen_knew = arg_parser.get_int("s_knew"); auto seqlen_kpads = arg_parser.get_int_vec("s_kpad"); + auto seqlen_qpads = arg_parser.get_int_vec("s_qpad"); + auto q_eff_lens_per_batch = arg_parser.get_int_vec("q_eff_lens"); + auto kv_eff_lens_per_batch = arg_parser.get_int_vec("kv_eff_lens"); ck_tile::index_t rotary_dim = arg_parser.get_int("rotary_dim"); bool i_perm = arg_parser.get_bool("iperm"); bool o_perm = arg_parser.get_bool("operm"); @@ -174,7 +189,10 @@ auto run(const ck_tile::ArgParser& arg_parser) hdim_q, hdim_v, seqlen_knew, + seqlen_qpads, seqlen_kpads, + q_eff_lens_per_batch, + kv_eff_lens_per_batch, rotary_dim, i_perm, o_perm, diff --git a/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp b/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp index 569c98a458..7ddb65a2db 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp @@ -52,7 +52,16 @@ auto parse_cmd_args(int argc, char* argv[]) -> std::pair get_query_shape() const @@ -172,6 +183,8 @@ struct Problem mask_info mask; TensorLayout input_layout; TensorLayout output_layout; + std::vector q_eff_lens; + std::vector kv_eff_lens; }; struct RunConfig @@ -326,8 +339,10 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) q_buf.ToDevice(q.data()); k_buf.ToDevice(k.data()); v_buf.ToDevice(v.data()); + // Ensure output buffer is zero-initialized so padded regions compare cleanly + o_buf.SetZero(); - ck_tile::fmha_fwd_v3_args args; + ck_tile::fmha_fwd_v3_args args{}; args.data_type = problem.data_type; args.batch = problem.batch; @@ -380,6 +395,60 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) : problem.seqlen_q * problem.hdim; args.batch_stride_o = problem.seqlen_q * problem.nhead_q * problem.hdim; + // Optional cumulative seqlen overrides (exclude PAD) + const bool has_varlen_q = !problem.q_eff_lens.empty() && problem.q_eff_lens[0] != -1; + const bool has_varlen_k = !problem.kv_eff_lens.empty() && problem.kv_eff_lens[0] != -1; + + auto make_effective_vec = [&](const std::vector& opt_vec, ck_tile::index_t fallback) { + std::vector eff; + if(!opt_vec.empty() && opt_vec[0] != -1) + { + eff.assign(opt_vec.begin(), opt_vec.end()); + if(eff.size() < static_cast(problem.batch)) + { + eff.resize(problem.batch, eff.back()); + } + } + else + { + eff.assign(problem.batch, fallback); + } + return eff; + }; + + const auto eff_q_vec = make_effective_vec(problem.q_eff_lens, problem.seqlen_q); + const auto eff_kv_vec = make_effective_vec(problem.kv_eff_lens, problem.seqlen_k); + + // Calculate cumulative sums for kernel arguments if varlen is used + std::vector cuq_cum, cukv_cum; + auto calculate_cumulative = [&](const std::vector& per_batch_vec, + std::vector& cum_vec) { + cum_vec.resize(per_batch_vec.size() + 1); + cum_vec[0] = 0; + for(std::size_t i = 0; i < per_batch_vec.size(); ++i) + cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i]; + }; + + if(has_varlen_q) + { + calculate_cumulative(eff_q_vec, cuq_cum); + } + if(has_varlen_k) + { + calculate_cumulative(eff_kv_vec, cukv_cum); + } + + ck_tile::DeviceMem cuq_buf(!cuq_cum.empty() ? cuq_cum.size() * sizeof(ck_tile::index_t) : 0); + ck_tile::DeviceMem cukv_buf(!cukv_cum.empty() ? cukv_cum.size() * sizeof(ck_tile::index_t) : 0); + cuq_buf.ToDevice(!cuq_cum.empty() ? cuq_cum.data() : nullptr); + cukv_buf.ToDevice(!cukv_cum.empty() ? cukv_cum.data() : nullptr); + args.cu_seqlen_q_ptr = + !cuq_cum.empty() ? reinterpret_cast(cuq_buf.GetDeviceBuffer()) + : nullptr; + args.cu_seqlen_kv_ptr = + !cukv_cum.empty() ? reinterpret_cast(cukv_buf.GetDeviceBuffer()) + : nullptr; + ck_tile::stream_config stream_config{nullptr, true, /*log_level=*/0, @@ -442,15 +511,72 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) o_ref = o_ref.transpose({0, 2, 1, 3}); } - host::fmha_fwd(q, - k, - v, - problem.mask, - o_ref, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::scales{problem.softmax_scale}); + // If variable lengths are provided, compute per-batch references + // with the effective lengths; else compute a single full reference. + if(has_varlen_q || has_varlen_k) + { + // Variable-length aware verification: zero-fill padded region and only compute valid part. + o_ref.SetZero(); + + for(int b = 0; b < problem.batch; ++b) + { + const ck_tile::index_t seqlen_q_eff = eff_q_vec[b]; + const ck_tile::index_t seqlen_kv_eff = eff_kv_vec[b]; + + if(seqlen_q_eff <= 0 || seqlen_kv_eff <= 0) + continue; + + // Slice current batch from inputs (bshd) and build single-batch tensors + ck_tile::HostTensor q_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim}); + ck_tile::HostTensor k_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); + ck_tile::HostTensor v_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); + ck_tile::HostTensor o_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim}); + + // Copy effective region + q_b.ForEach([&](auto& self, auto idx) { + // idx: [0, s, h, d] + self(idx) = q(b, idx[1], idx[2], idx[3]); + }); + k_b.ForEach([&](auto& self, auto idx) { self(idx) = k(b, idx[1], idx[2], idx[3]); }); + v_b.ForEach([&](auto& self, auto idx) { self(idx) = v(b, idx[1], idx[2], idx[3]); }); + + // Compute reference for this batch segment (host::fmha_fwd expects bshd tensors) + host::fmha_fwd(q_b, + k_b, + v_b, + problem.mask, + o_b, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::scales{problem.softmax_scale}); + + // Scatter into o_ref's bshd descriptor memory + for(int s = 0; s < seqlen_q_eff; ++s) + { + for(int h = 0; h < problem.nhead_q; ++h) + { + for(int d = 0; d < problem.hdim; ++d) + { + o_ref(b, s, h, d) = o_b(0, s, h, d); + } + } + } + } + } + else + { + // No varlen override: compute the full reference once + host::fmha_fwd(q, + k, + v, + problem.mask, + o_ref, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::scales{problem.softmax_scale}); + } ck_tile::HostTensor o(problem.get_output_shape()); o_buf.FromDevice(o.data()); diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index c41e48e6aa..f5dd42a6bd 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -162,11 +162,20 @@ struct fmha_fwd_args void* lse_ptr; void* o_ptr; + // Optional cumulative sequence length arrays + // Batch mode: cu_seqlen_* override effective per-batch lengths (exclude PAD) + const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] + const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] + const void* seqstart_q_ptr; const void* seqstart_k_ptr; const void* seqlen_k_ptr; // only used if both 'seqstart_q_ptr' & 'seqstart_k_ptr' are not nullptr + // Group mode: seqstart_padded_* provide physical starts including PAD (optional) + const void* seqstart_padded_q_ptr = nullptr; // [batch+1] + const void* seqstart_padded_k_ptr = nullptr; // [batch+1] + ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; ck_tile::index_t batch; @@ -554,7 +563,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.min_seqlen_q, args.p_drop, args.s_randval, - args.drop_seed_offset); + args.drop_seed_offset, + args.seqstart_padded_q_ptr, + args.seqstart_padded_k_ptr); } else { // create batch mode kernel arguments @@ -600,7 +611,9 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset); + args.drop_seed_offset, + args.cu_seqlen_q_ptr, + args.cu_seqlen_kv_ptr); } }(); diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 43f484fe14..cb5827975e 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -151,7 +151,10 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t seqlen_knew, + std::vector seqlen_qpads, std::vector seqlen_kpads, + std::vector q_eff_lens_per_batch, + std::vector kv_eff_lens_per_batch, ck_tile::index_t rotary_dim, bool i_perm, bool o_perm, @@ -362,6 +365,44 @@ fwd_result fmha_fwd_run(mode_enum mode, const auto seqstart_k_host = to_seqstarts(seqlen_ks); const auto seqstart_k_with_padding_host = to_seqstarts(seqlen_kpads); + // Optional padded Q seqstarts (group-mode only) + std::vector seqstart_q_with_padding_host; + if(mode == mode_enum::group && !seqlen_qpads.empty() && seqlen_qpads[0] != -1) + { + if(seqlen_qpads.size() < static_cast(batch)) + { + seqlen_qpads.resize(batch, seqlen_qpads.back()); + } + if(seqlen_qpads.size() == static_cast(batch)) + { + seqstart_q_with_padding_host = to_seqstarts( + ck_tile::span(seqlen_qpads.data(), seqlen_qpads.size())); + } + } + + // Optional batch-mode cumulative seqlen overrides + std::vector cuq_cum, cukv_cum; + if(mode == mode_enum::batch) + { + auto calculate_cumulative = [&](std::vector& per_batch_vec, + std::vector& cum_vec) { + if(!per_batch_vec.empty() && per_batch_vec[0] != -1) + { + if(per_batch_vec.size() < static_cast(batch)) + { + per_batch_vec.resize(batch, per_batch_vec.back()); + } + cum_vec.resize(batch + 1); + cum_vec[0] = 0; + for(int i = 0; i < batch; ++i) + cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i]; + } + }; + + calculate_cumulative(q_eff_lens_per_batch, cuq_cum); + calculate_cumulative(kv_eff_lens_per_batch, cukv_cum); + } + using TypeConfig = FmhaFwdTypeConfig; using QDataType = typename TypeConfig::QDataType; @@ -445,8 +486,15 @@ fwd_result fmha_fwd_run(mode_enum mode, // host memory for storing all the tensor elements const ck_tile::index_t shape_batch = (mode == mode_enum::batch ? batch : 1); - const ck_tile::index_t shape_seqlen_q = + // logical(unpadded) total seqlen_q for group; batch uses fixed seqlen + const ck_tile::index_t shape_seqlen_q_lse = (mode == mode_enum::batch ? seqlen_qs[0] : seqstart_q_host.back()); + // physical(padded) total seqlen_q for group when s_qpad is provided; else use logical + const ck_tile::index_t shape_seqlen_q = + (mode == mode_enum::batch + ? seqlen_qs[0] + : (seqstart_q_with_padding_host.empty() ? seqstart_q_host.back() + : seqstart_q_with_padding_host.back())); const ck_tile::index_t shape_seqlen_k = (mode == mode_enum::batch ? seqlen_ks[0] : (seqlen_kpads[0] < 0 ? seqstart_k_host.back() @@ -504,7 +552,7 @@ fwd_result fmha_fwd_run(mode_enum mode, // batch mode of lse data layout is [batch, nhead, seqlen_q] // group mode of lse data layout is [nhead, total_seqlen_q] ck_tile::HostTensor lse_host( - lse ? std::array{shape_batch, nhead, shape_seqlen_q} + lse ? std::array{shape_batch, nhead, shape_seqlen_q_lse} : std::array{1, 1, 1} /* dummy shape for simplifying code */); ck_tile::HostTensor o_host( @@ -602,6 +650,16 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t)); ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem seqstart_q_padded_buf(seqstart_q_with_padding_host.empty() + ? 0 + : seqstart_q_with_padding_host.size() * + sizeof(int32_t)); + ck_tile::DeviceMem seqstart_k_padded_buf( + seqlen_kpads[0] < 0 ? 0 : seqstart_k_with_padding_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem cu_seqlen_q_buf(cuq_cum.empty() ? 0 + : cuq_cum.size() * sizeof(ck_tile::index_t)); + ck_tile::DeviceMem cu_seqlen_kv_buf( + cukv_cum.empty() ? 0 : cukv_cum.size() * sizeof(ck_tile::index_t)); ck_tile::DeviceMem seqlen_k_buf((mode == mode_enum::batch && use_kvcache) || 0 <= seqlen_kpads[0] ? seqlen_ks.size() * sizeof(int32_t) @@ -693,8 +751,14 @@ fwd_result fmha_fwd_run(mode_enum mode, vnew_buf.ToDevice(vnew_host.data()); bias_buf.ToDevice(bias_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); - seqstart_k.ToDevice(seqlen_kpads[0] < 0 ? seqstart_k_host.data() - : seqstart_k_with_padding_host.data()); + // Keep logical starts in seqstart_k; pass padded K via separate pointer + seqstart_k.ToDevice(seqstart_k_host.data()); + seqstart_q_padded_buf.ToDevice( + seqstart_q_with_padding_host.empty() ? nullptr : seqstart_q_with_padding_host.data()); + seqstart_k_padded_buf.ToDevice(seqlen_kpads[0] < 0 ? nullptr + : seqstart_k_with_padding_host.data()); + cu_seqlen_q_buf.ToDevice(cuq_cum.empty() ? nullptr : cuq_cum.data()); + cu_seqlen_kv_buf.ToDevice(cukv_cum.empty() ? nullptr : cukv_cum.data()); seqlen_k_buf.ToDevice((mode == mode_enum::batch && use_kvcache) || 0 <= seqlen_kpads[0] ? seqlen_ks.data() : nullptr); @@ -830,8 +894,8 @@ fwd_result fmha_fwd_run(mode_enum mode, const ck_tile::index_t nhead_stride_bias = (i_perm ? 0 * shape_seqlen_q * max_seqlen_k : 0 * max_seqlen_k); const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t nhead_stride_lse = shape_seqlen_q; - const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q); + const ck_tile::index_t nhead_stride_lse = shape_seqlen_q_lse; + const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q_lse); const ck_tile::index_t nhead_stride_o_acc = (num_splits * shape_seqlen_q * hdim_v); const ck_tile::index_t nhead_stride_o = (o_perm ? shape_seqlen_q * hdim_v : hdim_v); // setup batch_stride_* arguments @@ -846,8 +910,8 @@ fwd_result fmha_fwd_run(mode_enum mode, const ck_tile::index_t batch_stride_vnew = (nhead_k * hdim_v * seqlen_knew); const ck_tile::index_t batch_stride_bias = (0 * nhead * shape_seqlen_q * max_seqlen_k); const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q); - const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q); + const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q_lse); + const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q_lse); const ck_tile::index_t batch_stride_o_acc = (nhead * num_splits * shape_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_o = (nhead * shape_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_block_table = (max_num_page_blocks / batch); @@ -961,6 +1025,29 @@ fwd_result fmha_fwd_run(mode_enum mode, { args.drop_seed_offset = std::make_pair(drop_seed, drop_offset); } + + // Group-mode: optional physical padded starts for Q/K + if(mode == mode_enum::group) + { + args.seqstart_padded_q_ptr = (seqstart_q_with_padding_host.empty() + ? nullptr + : seqstart_q_padded_buf.GetDeviceBuffer()); + args.seqstart_padded_k_ptr = + (seqlen_kpads[0] < 0 ? nullptr : seqstart_k_padded_buf.GetDeviceBuffer()); + } + + // Batch-mode: optional cumulative effective seqlen overrides + if(mode == mode_enum::batch) + { + args.cu_seqlen_q_ptr = cuq_cum.empty() + ? nullptr + : reinterpret_cast( + cu_seqlen_q_buf.GetDeviceBuffer()); + args.cu_seqlen_kv_ptr = cukv_cum.empty() + ? nullptr + : reinterpret_cast( + cu_seqlen_kv_buf.GetDeviceBuffer()); + } } else if constexpr(std::is_same_v>) { @@ -1167,15 +1254,29 @@ fwd_result fmha_fwd_run(mode_enum mode, for(ck_tile::index_t wb = 0; wb < batch; ++wb) { - const ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb]; - const ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb]; + ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb]; + ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb]; + if(mode == mode_enum::batch) + { + if(!cuq_cum.empty()) + { + real_seqlen_q = cuq_cum[wb + 1] - cuq_cum[wb]; + } + if(!cukv_cum.empty()) + { + real_seqlen_k = cukv_cum[wb + 1] - cukv_cum[wb]; + } + } // adjust matrix index according to the mode const ck_tile::index_t b_idx = (mode == mode_enum::batch ? wb : 0); const ck_tile::index_t cache_b_idx = (use_cache_batch_idx ? cache_batch_idx_host(b_idx) : b_idx); const ck_tile::index_t query_offset = - (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]); + (mode == mode_enum::batch + ? 0 + : (seqstart_q_with_padding_host.empty() ? seqstart_q_host[wb] + : seqstart_q_with_padding_host[wb])); const ck_tile::index_t key_offset = (mode == mode_enum::batch ? 0 @@ -1538,8 +1639,10 @@ fwd_result fmha_fwd_run(mode_enum mode, if(lse) { ck_tile::HostTensor lse_host_result({nhead, real_seqlen_q}); + const ck_tile::index_t query_offset_lse = + (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]); lse_host_result.ForEach([&](auto& self, auto idx) { - self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset); + self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset_lse); }); cur_pass = ck_tile::check_err(lse_host_result, diff --git a/example/ck_tile/01_fmha/fmha_fwd_v3.hpp b/example/ck_tile/01_fmha/fmha_fwd_v3.hpp index 10cb5149a4..4bd1d1a367 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_v3.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_v3.hpp @@ -56,6 +56,11 @@ struct fmha_fwd_v3_args index_t stride_o; index_t nhead_stride_o; index_t batch_stride_o; + + // Optional batch-mode cumulative seqlen overrides (exclude PAD) + // If provided, they override per-batch effective lengths to skip tail padding. + const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] + const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] }; std::ostream& operator<<(std::ostream& stream, const fmha_fwd_v3_args::data_type_enum& data_type); diff --git a/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp b/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp index e0fbad39a5..194675f962 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp @@ -158,7 +158,9 @@ float fmha_fwd_v3_kernel_launch(const fmha_fwd_v3_args& args, const stream_confi args.window_size_left, args.window_size_right, args.mask_type, - remap_opt); + remap_opt, + args.cu_seqlen_q_ptr, + args.cu_seqlen_kv_ptr); dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, args.hdim_v); constexpr dim3 blocks = Kernel::BlockSize(); diff --git a/example/ck_tile/01_fmha/script/benchmark_fwd.sh b/example/ck_tile/01_fmha/script/benchmark_fwd.sh index 88c16cceb6..31ad800039 100755 --- a/example/ck_tile/01_fmha/script/benchmark_fwd.sh +++ b/example/ck_tile/01_fmha/script/benchmark_fwd.sh @@ -18,3 +18,36 @@ $EXE -prec=$prec -b=1 -h=$nhead -d=$hdim -s=16384 -iperm=$perm -operm=$perm -kn done done done + +#Padding Benchmarks: batch mode (baseline vs low/med/high pad) +prec="fp16" +base_batch_args="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID" + +# baseline (no pad) +$EXE $base_batch_args + +# low pad (≈90–95% effective) +$EXE $base_batch_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 + +# medium pad (≈60–75% effective) +$EXE $base_batch_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 + +# high pad (≈30–40% effective) +$EXE $base_batch_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 + +# Padding Benchmarks: group mode (baseline vs low/med/high physical pad) +seqlens_q="1024,768,512,256" +seqlens_k="1024,768,512,256" +base_group_args="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID" + +# baseline (no physical pad) +$EXE $base_group_args + +# low physical pad +$EXE $base_group_args -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320 + +# medium physical pad +$EXE $base_group_args -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384 + +# high physical pad +$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 diff --git a/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh b/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh index b847e85398..a3f7d68eb3 100755 --- a/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh +++ b/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh @@ -23,3 +23,20 @@ done done done done + +# Padding benchmark comparisons for v3 (batch mode only) +# ==== V3 Padding Benchmarks: batch mode (baseline vs low/med/high pad) ==== +prec="fp16" +base_v3_args="-prec=$prec -b=4 -h=16 -d=128 -s=1024 -mask=0 -iperm=0 -operm=0 -v=$VALID" + +# baseline (no pad) +$EXE $base_v3_args + +# low pad (≈90–95% effective) +$EXE $base_v3_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 + +# medium pad (≈60–75% effective) +$EXE $base_v3_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 + +# high pad (≈30–40% effective) +$EXE $base_v3_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 diff --git a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh index afd0c728c6..fca6b8d0cd 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh @@ -137,9 +137,118 @@ run_fp16_appendkv_tests() { done ; done ; done } +run_padding_smoke_tests() { + # Padding-only smoke tests for batch/group mode using COMMON_ARGS + local prec="fp16" + + # Batch mode: padding via effective lengths (exclude PAD) + # Use lse=1 to select a non-trload kernel and avoid overly strict tolerance mismatches + local base_batch="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=1 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS" + # low pad (≈90–95% effective) + $EXE $base_batch -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 + # medium pad (≈60–75% effective) + $EXE $base_batch -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 + # high pad (≈30–40% effective) + $EXE $base_batch -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 + + # Group mode: padding via physical stride along seqlen + local seqlens_q="1024,768,512,256" + local seqlens_k="1024,768,512,256" + local base_group="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS" + # low physical pad + $EXE $base_group -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320 + # medium physical pad + $EXE $base_group -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384 + # high physical pad + $EXE $base_group -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 +} + +run_padding_basic_boundary_tests() { + # Basic padding and boundary tests (reference: smoke_test_fwd_pad.sh) + local prec + local perm + + # Group mode: Q&K padded with per-batch different strides + for prec in fp16 bf16 ; do + for perm in 0 1 ; do + $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=16 -d_v=32 \ + -s=55 -s_k=256 -s_qpad=64,60 -s_kpad=272,260 \ + -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \ + -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS + done + done + + # slightly larger, uneven padding strides + for prec in fp16 bf16 ; do + for perm in 0 1 ; do + $EXE -prec=$prec -mode=1 -b=3 -h=2 -h_k=1 -d=64 -d_v=64 \ + -s=50,60,40 -s_k=128,256,192 -s_qpad=64,64,64 -s_kpad=160,288,224 \ + -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ + -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS + done + done + + # only K padded; Q unpadded + for prec in fp16 bf16 ; do + for perm in 0 1 ; do + $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 \ + -s=55 -s_k=256 -s_kpad=272,260 \ + -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ + -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS + done + done + + # use cu_seqlen overrides to skip tail PAD + for prec in fp16 bf16 ; do + for perm in 0 1 ; do + $EXE -prec=$prec -mode=0 -b=4 -h=8 -h_k=8 -d=128 -s=3 -s_k=3 \ + -q_eff_lens=1,2,1,2 -kv_eff_lens=1,2,1,2 \ + -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ + -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS + + $EXE -prec=$prec -mode=0 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 -s=64 -s_k=256 \ + -q_eff_lens=55,60 -kv_eff_lens=200,256 \ + -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \ + -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS + done + done + + # no padding (equal), mixed Q/KV, all len=1 + for prec in fp16 bf16 ; do + $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ + -q_eff_lens=128,128,128,128 -kv_eff_lens=128,128,128,128 \ + -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS + + $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ + -q_eff_lens=10,20,30,40 -kv_eff_lens=40,30,20,10 \ + -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS + + $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ + -q_eff_lens=1,1,1,1 -kv_eff_lens=1,1,1,1 \ + -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS + done + + # highly variable logical lengths + for prec in fp16 bf16 ; do + $EXE -prec=$prec -mode=1 -b=4 -h=4 -d=32 \ + -s=1,127,3,65 -s_k=1,127,3,65 -s_kpad=128 \ + -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS + done + + # GQA + Alibi + Causal mask (keep vlayout row-major for fp16/bf16 + for prec in fp16 bf16 ; do + $EXE -prec=$prec -mode=1 -b=2 -h=16 -h_k=4 -d=128 \ + -s=256,129 -s_k=256,129 -s_kpad=256 \ + -bias=a -mask=t -lse=1 -iperm=0 -operm=0 -vlayout=r \ + -kname=$KNAME $COMMON_ARGS + done +} + set -x run_fp16_bf16_tests +run_padding_smoke_tests +run_padding_basic_boundary_tests run_fp8_tests run_fp8bf16_tests run_fp8fp32_tests diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 58fdad149a..3f417bc125 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -291,6 +291,11 @@ struct FmhaFwdKernel ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; ck_tile::index_t batch_stride_o; + + // Optional cumulative sequence length pointers for batch mode + // If provided, they override seqlen_q / seqlen_k per-batch to skip tail padding. + const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // cumulative, length without PAD + const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // cumulative, length without PAD }; struct FmhaFwdGroupModeKargs @@ -310,6 +315,11 @@ struct FmhaFwdKernel const int32_t* seqstart_q_ptr; const int32_t* seqstart_k_ptr; const int32_t* seqlen_k_ptr; + + // Optional cumulative padded sequence starts (including PAD tokens) + // Used solely to compute memory offsets when sequences are physically padded. + const int32_t* seqstart_padded_q_ptr = nullptr; + const int32_t* seqstart_padded_k_ptr = nullptr; }; using Kargs = std::conditional_t; @@ -460,6 +470,105 @@ struct FmhaFwdKernel return kargs; } + // Overload: Batch mode with optional cu_seqlen pointers (unpadded cumulative lengths) + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargsImpl(const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* bias_ptr, + void* rand_val_ptr, + void* lse_ptr, + void* o_ptr, + ck_tile::index_t seqlen_q, + ck_tile::index_t seqlen_k, + ck_tile::index_t hdim_q, + ck_tile::index_t hdim_v, + ck_tile::index_t num_head_q, + ck_tile::index_t nhead_ratio_qk, + float scale_s, + float scale_p, + float scale_o, + float logits_soft_cap, + ck_tile::index_t stride_q, + ck_tile::index_t stride_k, + ck_tile::index_t stride_v, + ck_tile::index_t stride_bias, + ck_tile::index_t stride_randval, + ck_tile::index_t stride_o, + ck_tile::index_t nhead_stride_q, + ck_tile::index_t nhead_stride_k, + ck_tile::index_t nhead_stride_v, + ck_tile::index_t nhead_stride_bias, + ck_tile::index_t nhead_stride_randval, + ck_tile::index_t nhead_stride_lse, + ck_tile::index_t nhead_stride_o, + ck_tile::index_t batch_stride_q, + ck_tile::index_t batch_stride_k, + ck_tile::index_t batch_stride_v, + ck_tile::index_t batch_stride_bias, + ck_tile::index_t batch_stride_randval, + ck_tile::index_t batch_stride_lse, + ck_tile::index_t batch_stride_o, + ck_tile::index_t window_size_left, + ck_tile::index_t window_size_right, + ck_tile::index_t mask_type, + float p_drop, + bool s_randval, + std::variant, std::pair> + drop_seed_offset, + const ck_tile::index_t* cu_seqlen_q_ptr, + const ck_tile::index_t* cu_seqlen_kv_ptr) + { + auto kargs = MakeKargsImpl(q_ptr, + k_ptr, + v_ptr, + bias_ptr, + rand_val_ptr, + lse_ptr, + o_ptr, + seqlen_q, + seqlen_k, + hdim_q, + hdim_v, + num_head_q, + nhead_ratio_qk, + scale_s, + scale_p, + scale_o, + logits_soft_cap, + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_bias, + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + batch_stride_bias, + batch_stride_randval, + batch_stride_lse, + batch_stride_o, + window_size_left, + window_size_right, + mask_type, + p_drop, + s_randval, + drop_seed_offset); + + kargs.cu_seqlen_q_ptr = cu_seqlen_q_ptr; + kargs.cu_seqlen_kv_ptr = cu_seqlen_kv_ptr; + return kargs; + } + // std::variant<> can't take in a list initializer, overload for backward compatibility template CK_TILE_HOST static constexpr std::enable_if_t @@ -781,6 +890,95 @@ struct FmhaFwdKernel return kargs; } + // Overload: Group mode with optional padded seqstarts for memory offsets + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargsImpl(const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + const void* bias_ptr, + void* rand_val_ptr, + void* lse_ptr, + void* o_ptr, + const void* seqstart_q_ptr, + const void* seqstart_k_ptr, + const void* seqlen_k_ptr, + ck_tile::index_t hdim_q, + ck_tile::index_t hdim_v, + ck_tile::index_t num_head_q, + ck_tile::index_t nhead_ratio_qk, + float scale_s, + float scale_p, + float scale_o, + float logits_soft_cap, + ck_tile::index_t stride_q, + ck_tile::index_t stride_k, + ck_tile::index_t stride_v, + ck_tile::index_t stride_bias, + ck_tile::index_t stride_randval, + ck_tile::index_t stride_o, + ck_tile::index_t nhead_stride_q, + ck_tile::index_t nhead_stride_k, + ck_tile::index_t nhead_stride_v, + ck_tile::index_t nhead_stride_bias, + ck_tile::index_t nhead_stride_randval, + ck_tile::index_t nhead_stride_lse, + ck_tile::index_t nhead_stride_o, + ck_tile::index_t window_size_left, + ck_tile::index_t window_size_right, + ck_tile::index_t mask_type, + ck_tile::index_t min_seqlen_q, + float p_drop, + bool s_randval, + std::variant, std::pair> + drop_seed_offset, + const void* seqstart_padded_q_ptr, + const void* seqstart_padded_k_ptr) + { + auto kargs = MakeKargsImpl(q_ptr, + k_ptr, + v_ptr, + bias_ptr, + rand_val_ptr, + lse_ptr, + o_ptr, + seqstart_q_ptr, + seqstart_k_ptr, + seqlen_k_ptr, + hdim_q, + hdim_v, + num_head_q, + nhead_ratio_qk, + scale_s, + scale_p, + scale_o, + logits_soft_cap, + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_bias, + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + window_size_left, + window_size_right, + mask_type, + min_seqlen_q, + p_drop, + s_randval, + drop_seed_offset); + + kargs.seqstart_padded_q_ptr = reinterpret_cast(seqstart_padded_q_ptr); + kargs.seqstart_padded_k_ptr = reinterpret_cast(seqstart_padded_k_ptr); + return kargs; + } + // std::variant<> can't take in a list initializer, overload for backward compatibility template CK_TILE_HOST static constexpr std::enable_if_t @@ -1073,35 +1271,44 @@ struct FmhaFwdKernel if constexpr(kIsGroupMode) { - // get starting offset for each batch - const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; + // logical and physical (padded) starts + const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; - batch_offset_q = query_start * kargs.stride_q; - batch_offset_k = key_start * kargs.stride_k; + const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr + ? kargs.seqstart_padded_q_ptr[i_batch] + : query_start_unpadded; + const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr + ? kargs.seqstart_padded_k_ptr[i_batch] + : key_start_unpadded; + + // DRAM base offsets use physical padded starts + batch_offset_q = query_start_padded * kargs.stride_q; + batch_offset_k = key_start_padded * kargs.stride_k; if constexpr(std::is_same_v) { - batch_offset_v = key_start * kargs.stride_v; + batch_offset_v = key_start_padded * kargs.stride_v; } else { - batch_offset_v = key_start; + batch_offset_v = key_start_padded; } if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { - batch_offset_bias = query_start * kargs.stride_bias; + batch_offset_bias = query_start_padded * kargs.stride_bias; } if constexpr(kStoreLSE) { - batch_offset_lse = query_start; + // LSE stays indexed by unpadded starts + batch_offset_lse = query_start_unpadded; } if constexpr(kHasDropout) { - batch_offset_randval = query_start * kargs.stride_randval; + batch_offset_randval = query_start_padded * kargs.stride_randval; } - batch_offset_o = query_start * kargs.stride_o; + batch_offset_o = query_start_padded * kargs.stride_o; - // get real # queries & # keys under group mode + // real logical lengths (exclude PAD) const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0]; @@ -1113,8 +1320,7 @@ struct FmhaFwdKernel } } - // # of required blocks is different in each groups, terminate unnecessary blocks - // earlier + // terminate unnecessary blocks earlier if(kargs.seqlen_q <= i_m0) { return; @@ -1150,6 +1356,18 @@ struct FmhaFwdKernel static_cast(i_batch) * kargs.batch_stride_randval; } batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; + + // If cumulative seqlen pointers are provided, override per-batch effective lengths + if(kargs.cu_seqlen_q_ptr != nullptr) + { + kargs.seqlen_q = + kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; + } + if(kargs.cu_seqlen_kv_ptr != nullptr) + { + kargs.seqlen_k = + kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; + } } // for simplicity, batch stride we just modify the pointer @@ -1548,26 +1766,35 @@ struct FmhaFwdKernel if constexpr(kIsGroupMode) { // get starting offset for each batch - const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; + const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; - batch_offset_q = query_start * kargs.stride_q; - batch_offset_k = key_start * kargs.stride_k; + const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr + ? kargs.seqstart_padded_q_ptr[i_batch] + : query_start_unpadded; + const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr + ? kargs.seqstart_padded_k_ptr[i_batch] + : key_start_unpadded; + + batch_offset_q = query_start_padded * kargs.stride_q; + batch_offset_k = key_start_padded * kargs.stride_k; if constexpr(std::is_same_v) { - batch_offset_v = key_start * kargs.stride_v; + batch_offset_v = key_start_padded * kargs.stride_v; } else { - batch_offset_v = key_start; + // col-major V: offset along seqlen dimension is scalar index + batch_offset_v = key_start_padded; } if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { - batch_offset_bias = query_start * kargs.stride_bias; + batch_offset_bias = query_start_padded * kargs.stride_bias; } - batch_offset_lse = query_start; - batch_offset_o = query_start * kargs.stride_o; + // LSE layout is [nhead, total_seqlen], index by unpadded start + batch_offset_lse = query_start_unpadded; + batch_offset_o = query_start_padded * kargs.stride_o; // get real # queries & # keys under group mode kargs.seqlen_q = kargs.seqstart_q_ptr[i_batch + 1] - kargs.seqstart_q_ptr[i_batch]; @@ -1605,6 +1832,18 @@ struct FmhaFwdKernel batch_offset_bias = static_cast(i_batch) * kargs.batch_stride_bias; } + + // If cumulative seqlen pointers are provided, override per-batch effective lengths + if(kargs.cu_seqlen_q_ptr != nullptr) + { + kargs.seqlen_q = + kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; + } + if(kargs.cu_seqlen_kv_ptr != nullptr) + { + kargs.seqlen_k = + kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; + } } // for simplicity, batch stride we just modify the pointer diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp index c5e5745817..52b9da40b8 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp @@ -100,6 +100,11 @@ struct FmhaFwdV3Kernel ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; ck_tile::index_t batch_stride_o; + + // Optional cumulative sequence length pointers for batch mode + // If provided, they override seqlen_q / seqlen_k per-batch to skip tail padding. + const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] + const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] }; struct FmhaFwdGroupModeKargs @@ -110,6 +115,11 @@ struct FmhaFwdV3Kernel const int32_t* seqstart_q_ptr; const int32_t* seqstart_k_ptr; const int32_t* seqlen_k_ptr; + + // Optional cumulative padded sequence starts (including PAD tokens) + // Used solely to compute memory offsets when sequences are physically padded. + const int32_t* seqstart_padded_q_ptr = nullptr; // [batch+1] + const int32_t* seqstart_padded_k_ptr = nullptr; // [batch+1] }; using Kargs = std::conditional_t; @@ -190,6 +200,78 @@ struct FmhaFwdV3Kernel return kargs; } + // Overload: Batch mode with optional cu_seqlen pointers + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargs(const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* lse_ptr, + void* o_ptr, + ck_tile::index_t seqlen_q, + ck_tile::index_t seqlen_k, + ck_tile::index_t hdim_q, + ck_tile::index_t hdim_v, + ck_tile::index_t num_head_q, + ck_tile::index_t nhead_ratio_qk, + float scale_s, + ck_tile::index_t stride_q, + ck_tile::index_t stride_k, + ck_tile::index_t stride_v, + ck_tile::index_t stride_o, + ck_tile::index_t nhead_stride_q, + ck_tile::index_t nhead_stride_k, + ck_tile::index_t nhead_stride_v, + ck_tile::index_t nhead_stride_lse, + ck_tile::index_t nhead_stride_o, + ck_tile::index_t batch_stride_q, + ck_tile::index_t batch_stride_k, + ck_tile::index_t batch_stride_v, + ck_tile::index_t batch_stride_lse, + ck_tile::index_t batch_stride_o, + ck_tile::index_t window_size_left, + ck_tile::index_t window_size_right, + ck_tile::index_t mask_type, + ck_tile::index_t remap_opt, + const ck_tile::index_t* cu_seqlen_q_ptr, + const ck_tile::index_t* cu_seqlen_kv_ptr) + { + auto kargs = MakeKargs(q_ptr, + k_ptr, + v_ptr, + lse_ptr, + o_ptr, + seqlen_q, + seqlen_k, + hdim_q, + hdim_v, + num_head_q, + nhead_ratio_qk, + scale_s, + stride_q, + stride_k, + stride_v, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + batch_stride_lse, + batch_stride_o, + window_size_left, + window_size_right, + mask_type, + remap_opt); + + kargs.cu_seqlen_q_ptr = cu_seqlen_q_ptr; + kargs.cu_seqlen_kv_ptr = cu_seqlen_kv_ptr; + return kargs; + } + template CK_TILE_HOST static constexpr std::enable_if_t MakeKargs(const void* q_ptr, @@ -260,6 +342,70 @@ struct FmhaFwdV3Kernel return kargs; } + // Overload: Group mode with optional padded seqstarts for memory offsets + template + CK_TILE_HOST static constexpr std::enable_if_t + MakeKargs(const void* q_ptr, + const void* k_ptr, + const void* v_ptr, + void* lse_ptr, + void* o_ptr, + const void* seqstart_q_ptr, + const void* seqstart_k_ptr, + const void* seqlen_k_ptr, + ck_tile::index_t hdim_q, + ck_tile::index_t hdim_v, + ck_tile::index_t num_head_q, + ck_tile::index_t nhead_ratio_qk, + float scale_s, + ck_tile::index_t stride_q, + ck_tile::index_t stride_k, + ck_tile::index_t stride_v, + ck_tile::index_t stride_o, + ck_tile::index_t nhead_stride_q, + ck_tile::index_t nhead_stride_k, + ck_tile::index_t nhead_stride_v, + ck_tile::index_t nhead_stride_lse, + ck_tile::index_t nhead_stride_o, + ck_tile::index_t window_size_left, + ck_tile::index_t window_size_right, + ck_tile::index_t mask_type, + ck_tile::index_t remap_opt, + const void* seqstart_padded_q_ptr, + const void* seqstart_padded_k_ptr) + { + auto kargs = MakeKargs(q_ptr, + k_ptr, + v_ptr, + lse_ptr, + o_ptr, + seqstart_q_ptr, + seqstart_k_ptr, + seqlen_k_ptr, + hdim_q, + hdim_v, + num_head_q, + nhead_ratio_qk, + scale_s, + stride_q, + stride_k, + stride_v, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_lse, + nhead_stride_o, + window_size_left, + window_size_right, + mask_type, + remap_opt); + + kargs.seqstart_padded_q_ptr = reinterpret_cast(seqstart_padded_q_ptr); + kargs.seqstart_padded_k_ptr = reinterpret_cast(seqstart_padded_k_ptr); + return kargs; + } + CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_, @@ -373,18 +519,26 @@ struct FmhaFwdV3Kernel if constexpr(kIsGroupMode) { // get starting offset for each batch - const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; + const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; - batch_offset_q = query_start * kargs.stride_q; - batch_offset_k = key_start * kargs.stride_k; - batch_offset_v = key_start * kargs.stride_v; + const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr + ? kargs.seqstart_padded_q_ptr[i_batch] + : query_start_unpadded; + const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr + ? kargs.seqstart_padded_k_ptr[i_batch] + : key_start_unpadded; + + batch_offset_q = query_start_padded * kargs.stride_q; + batch_offset_k = key_start_padded * kargs.stride_k; + batch_offset_v = key_start_padded * kargs.stride_v; if constexpr(kStoreLSE) { - batch_offset_lse = query_start; + // LSE layout is [nhead, total_seqlen], index by unpadded start + batch_offset_lse = query_start_unpadded; } - batch_offset_o = query_start * kargs.stride_o; + batch_offset_o = query_start_padded * kargs.stride_o; // get real # queries & # keys under group mode const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; @@ -417,6 +571,18 @@ struct FmhaFwdV3Kernel batch_offset_lse = static_cast(i_batch) * kargs.batch_stride_lse; } batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; + + // If cumulative seqlen pointers are provided, override per-batch effective lengths + if(kargs.cu_seqlen_q_ptr != nullptr) + { + kargs.seqlen_q = + kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; + } + if(kargs.cu_seqlen_kv_ptr != nullptr) + { + kargs.seqlen_k = + kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; + } } // for simplicity, batch stride we just modify the pointer diff --git a/test/ck_tile/fmha/test_fmha_fwd.inc b/test/ck_tile/fmha/test_fmha_fwd.inc index 08abd3358d..66d4e3dc21 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.inc +++ b/test/ck_tile/fmha/test_fmha_fwd.inc @@ -98,7 +98,10 @@ TEST_P(AllLong, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim perm, // i_perm perm, // o_perm @@ -160,7 +163,10 @@ TEST_P(HDimPadding, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim perm, // i_perm perm, // o_perm @@ -217,7 +223,10 @@ TEST_P(ElementwiseBias, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -273,7 +282,10 @@ TEST_P(Alibi, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim true, // i_perm true, // o_perm @@ -331,7 +343,10 @@ TEST_P(Dropout, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim false, // i_perm false, // o_perm @@ -391,7 +406,10 @@ TEST_P(PagedKV, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -457,7 +475,10 @@ TEST_P(SplitKV, Test) hdim_q, hdim_v, 0, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -529,7 +550,10 @@ TEST_P(AppendKV, Test) hdim_q, hdim_v, seqlen_knew, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm true, // o_perm @@ -599,7 +623,10 @@ TEST_P(AppendKVRoPE, Test) hdim_q, hdim_v, seqlen_knew, // seqlen_knew + {-1}, // seqlen_qpads {-1}, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch rotary_dim, // rotary_dim i_perm, // i_perm true, // o_perm @@ -623,3 +650,117 @@ TEST_P(AppendKVRoPE, Test) } #endif // CK_TILE_FMHA_FWD_APPENDKV_API + +// --------------------------------------------------------------- +// Additional padding tests (q/kv physical padding & effective len) +// --------------------------------------------------------------- + +// Simple batch-mode test with per-batch Q/KV padding strides and effective lengths +TEST(TestCkTileFmhaFwd, BatchModeQKvPadding) +{ + if constexpr(std::is_same_v) + { + GTEST_SKIP() << "Skip for fp8"; + } + const mode_enum mode = mode_enum::batch; + const int batch = 3; + const int nhead = 2; + const int nhead_k = -1; + const int seqlen_q = 128; + const int seqlen_k = 128; + const int hdim_q = 64; + const int hdim_v = 64; + const int seqlen_knew = 0; + const std::vector seqlen_qpads{}; + const std::vector seqlen_kpads{}; + const std::vector q_eff_lens{120, 128, 100}; + const std::vector kv_eff_lens{110, 128, 90}; + + auto result = fmha_fwd_run(mode, + batch, + nhead, + nhead_k, + {adjust_seqlen(seqlen_q)}, + {adjust_seqlen(seqlen_k)}, + hdim_q, + hdim_v, + seqlen_knew, // seqlen_knew + seqlen_qpads, // seqlen_qpads + seqlen_kpads, // seqlen_kpads + q_eff_lens, // q_eff_lens_per_batch + kv_eff_lens, // kv_eff_lens_per_batch + 0, // rotary_dim + true, // i_perm + true, // o_perm + 0, // scale_s + 0, // logits_soft_cap + def_is_v_rowmajor, + def_lse, // lse + 0, // page_block_size + false, // use_cache_batch_idx + "n", // bias_str + 0.0f, // p_drop + 0, // drop_seed + 0, // drop_offset + false, // drop_prefs + "0", // mask_str + QUANT_ARGS, + true, // is_rotary_interleaved + 1, // num_splits + COMMON_ARGS); + CHECK_RESULT(result); +} + +// Simple group-mode test with uniform seqlen but per-batch padding & effective lengths +TEST(TestCkTileFmhaFwd, GroupModeQKvPadding) +{ + if constexpr(std::is_same_v) + { + GTEST_SKIP() << "Skip for fp8"; + } + const mode_enum mode = mode_enum::group; + const int batch = 2; + const int nhead = 2; + const int nhead_k = -1; + const std::vector seqlen_q{96, 128}; // unpadded + const std::vector seqlen_k{96, 128}; // unpadded + const int hdim_q = 64; + const int hdim_v = 64; + const int seqlen_knew = 0; + const std::vector seqlen_qpads{128, 160}; + const std::vector seqlen_kpads{128, 160}; + + auto result = fmha_fwd_run(mode, + batch, + nhead, + nhead_k, + seqlen_q, + seqlen_k, + hdim_q, + hdim_v, + seqlen_knew, // seqlen_knew + seqlen_qpads, // seqlen_qpads + seqlen_kpads, // seqlen_kpads + {}, // q_eff_lens_per_batch + {}, // kv_eff_lens_per_batch + 0, // rotary_dim + true, // i_perm + true, // o_perm + 0, // scale_s + 0, // logits_soft_cap + def_is_v_rowmajor, + def_lse, // lse + 0, // page_block_size + false, // use_cache_batch_idx + "n", // bias_str + 0.0f, // p_drop + 0, // drop_seed + 0, // drop_offset + false, // drop_prefs + "0", // mask_str + QUANT_ARGS, + true, // is_rotary_interleaved + 1, // num_splits + COMMON_ARGS); + CHECK_RESULT(result); +} From 6cf3fdd21c502249767f814a087fbd9be88013eb Mon Sep 17 00:00:00 2001 From: Yi DING Date: Fri, 19 Sep 2025 21:45:02 +0800 Subject: [PATCH 07/12] [CK_TILE] FMHA BWD Fix Decode Accuracy (#2881) * [CK_TILE] FMHA BWD Fix Decode Accuracy * use s_waitcnt utils --- .../block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp index 8c8d2af486..6d90429407 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp @@ -489,7 +489,7 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR move_tile_window(k_dram_window, {kN0, 0}); async_load_tile(v_lds_write_window, v_dram_window); move_tile_window(v_dram_window, {kN0, 0}); - // __builtin_amdgcn_s_waitcnt(0); + s_waitcnt(); k_reg_tensor = load_tile(k_lds_read_window); v_reg_tensor = load_tile(v_lds_read_window); kt_reg_tensor = load_tile_transpose(kt_lds_read_window); @@ -636,7 +636,7 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR } }(); store_tile(bias_lds_write_window, dbias); - __builtin_amdgcn_s_waitcnt(3952); + s_waitcnt(); block_sync_lds(); auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); auto dbias_tile = make_static_distributed_tensor( @@ -664,7 +664,7 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR } store_tile(ds_lds_window, ds_gemm); } - __builtin_amdgcn_s_waitcnt(3952); + s_waitcnt(); block_sync_lds(); if constexpr(is_epilogue) { From 29446da1d57170a8bda47a452113ef7e44363a04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Fri, 19 Sep 2025 16:27:50 +0200 Subject: [PATCH 08/12] Disable bwd weight split-k autodeduce for single stage kernels (#2856) * Disable bwd weight split-k autodeduce for single stage kernels * update interface tests --------- Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> --- .../device/device_grouped_conv_bwd_weight.hpp | 2 + ...ice_grouped_conv_bwd_weight_multiple_d.hpp | 2 + ...e_grouped_conv_bwd_weight_explicit_xdl.hpp | 47 +++++++++++++--- ...onv_bwd_weight_multiple_d_xdl_cshuffle.hpp | 8 +++ ...e_grouped_conv_bwd_weight_xdl_cshuffle.hpp | 8 +++ ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 9 ++++ ...rouped_convnd_bwd_weight_interface_xdl.cpp | 53 ++++++++++--------- 7 files changed, 96 insertions(+), 33 deletions(-) diff --git a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp index 7296e4faaa..18223c78f7 100644 --- a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp +++ b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp @@ -11,6 +11,8 @@ namespace ck { namespace tensor_operation { namespace device { +#define DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS 1 + template ()) diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp index 934dc7ee8e..987a1e273a 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp @@ -671,6 +671,7 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -683,6 +684,7 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle grid_size); } else +#endif { k_batch_ = split_k; } @@ -939,6 +941,12 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp index b361409e38..22fc13bae4 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp @@ -553,6 +553,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle conv_ngchw_to_nhwgc_transformer.TransposeWeiStrides(e_g_k_c_xs_lengths, e_g_k_c_xs_strides); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -565,6 +566,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle grid_size); } else +#endif { k_batch_ = split_k; } @@ -934,6 +936,12 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index 8bf188be2e..735eebbdf6 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -524,6 +524,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN, gemmK; @@ -549,6 +550,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 } } else +#endif { k_batch_ = split_k; } @@ -1275,6 +1277,13 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif + const index_t GemmM = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I1); const index_t GemmN = arg.b_grid_desc_kbatch_k0_n_k1_.GetLength(I1); const index_t GemmK = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I0) * diff --git a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp index 2a9421fcd1..354d1fc23b 100644 --- a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp +++ b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp @@ -52,7 +52,7 @@ class TestGroupedConvndBwdWeight : public ::testing::Test // clang-format on ck::utils::conv::ConvParam conv_param; - std::vector split_ks{-1, 2}; + ck::index_t split_k_ = 2; template bool Run() @@ -96,30 +96,24 @@ class TestGroupedConvndBwdWeight : public ::testing::Test auto conv = GroupedConvBwdWeightDeviceInstance{}; - bool is_supported = true; - - for(const auto split_k : split_ks) - { - auto argument = conv.MakeArgument(nullptr, - nullptr, - nullptr, - input_lengths, - input_strides, - filter_lengths, - weights_strides, - output_lengths, - output_strides, - conv_filter_strides, - conv_filter_dilations, - input_left_pads, - input_right_pads, - PassThrough{}, - PassThrough{}, - PassThrough{}, - split_k); - is_supported &= conv.IsSupportedArgument(argument); - } - return is_supported; + auto argument = conv.MakeArgument(nullptr, + nullptr, + nullptr, + input_lengths, + input_strides, + filter_lengths, + weights_strides, + output_lengths, + output_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads, + PassThrough{}, + PassThrough{}, + PassThrough{}, + split_k_); + return conv.IsSupportedArgument(argument); } }; @@ -183,3 +177,12 @@ TYPED_TEST(TestGroupedConvndBwdWeightDefault, VectorLoadCheck) is_supported = this->template Run<2>(); EXPECT_FALSE(is_supported); } + +TYPED_TEST(TestGroupedConvndBwdWeightDefault, SingleStageAutoDeduce) +{ + // Supported version but with auto deduce and single stage + this->conv_param = {2, 2, 128, 128, 256, {1, 1}, {3, 3}, {1, 1}, {1, 1}, {0, 0}, {0, 0}}; + this->split_k_ = -1; + bool is_supported = this->template Run<2>(); + EXPECT_FALSE(is_supported); +} From b765fe78f37c85a9ca10c24fec6b7247a170034f Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 19 Sep 2025 08:15:02 -0700 Subject: [PATCH 09/12] =?UTF-8?q?Revert=20"[CK=5FTILE]=20Add=20sequence=20?= =?UTF-8?q?padding=20and=20variable=20length=20support=20in=20fmha=20(a?= =?UTF-8?q?=E2=80=A6"=20(#2883)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 86dd59cd01e41a4190bf2405a0fb0e89d9498b4c. --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 6 +- example/ck_tile/01_fmha/example_fmha_fwd.cpp | 20 +- .../ck_tile/01_fmha/example_fmha_fwd_v3.cpp | 148 +-------- example/ck_tile/01_fmha/fmha_fwd.hpp | 17 +- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 127 +------- example/ck_tile/01_fmha/fmha_fwd_v3.hpp | 5 - example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp | 4 +- .../ck_tile/01_fmha/script/benchmark_fwd.sh | 33 -- .../01_fmha/script/benchmark_fwd_v3.sh | 17 -- .../ck_tile/01_fmha/script/smoke_test_fwd.sh | 109 ------- .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 285 ++---------------- .../ops/fmha/kernel/fmha_fwd_v3_kernel.hpp | 180 +---------- test/ck_tile/fmha/test_fmha_fwd.inc | 141 --------- 13 files changed, 60 insertions(+), 1032 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index da0c9ca931..cfb96b7d53 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -259,11 +259,11 @@ class FmhaFwdApiTrait: def skcheck(self) -> str: if self.mode == 'group': return 'true/*group mode skpad always true*/' # group mode only generate spad/skpad == true if self.pipeline_tag == 'qr_async': - if self.skpad == 't' : return f'(a.cu_seqlen_kv_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0)' - else : return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)' + if self.skpad == 't' : return f'a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0' + else : return f'a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0' elif self.pipeline_tag in ['qr', 'qs']: if self.skpad == 't' : return f'true /*a.seqlen_k % {self.bn0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) - else : return f'(a.cu_seqlen_kv_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)' + else : return f'a.seqlen_k % {self.bn0} == 0' elif self.pipeline_tag == 'qr_async_trload': if self.skpad == 't' : return 'true' else: return 'true' diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp index 79fda6d564..91cb9f55be 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp @@ -33,10 +33,6 @@ auto create_args(int argc, char* argv[]) "0", "seqlen_k for new key/value, 0 means not to use this at all; " "-1 to choose s_knew in [1, s] randomly.") - .insert("s_qpad", - "-1", - "seqlen_q stride between 2 batches (group-mode optional).\n" - "Provide positive strides per-batch to simulate physical padding on Q.") .insert("s_kpad", "-1", "seqlen_k stride between 2 batches, currently used in group-mode only\n" @@ -111,15 +107,7 @@ auto create_args(int argc, char* argv[]) .insert("warmup", "5", "number of iterations before benchmark the kernel") .insert("repeat", "20", "number of iterations to benchmark the kernel") .insert("json", "0", "0: No Json, 1: Dump Results in Json format") - .insert("jsonfile", "fmha_fwd.json", "json file name to dump results") - .insert("q_eff_lens", - "", - "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n" - "Comma-separated list of length 'b'. If empty, no override.") - .insert("kv_eff_lens", - "", - "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n" - "Comma-separated list of length 'b'. If empty, no override."); + .insert("jsonfile", "fmha_fwd.json", "json file name to dump results"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); @@ -139,9 +127,6 @@ auto run(const ck_tile::ArgParser& arg_parser) ck_tile::index_t hdim_v = arg_parser.get_int("d_v"); ck_tile::index_t seqlen_knew = arg_parser.get_int("s_knew"); auto seqlen_kpads = arg_parser.get_int_vec("s_kpad"); - auto seqlen_qpads = arg_parser.get_int_vec("s_qpad"); - auto q_eff_lens_per_batch = arg_parser.get_int_vec("q_eff_lens"); - auto kv_eff_lens_per_batch = arg_parser.get_int_vec("kv_eff_lens"); ck_tile::index_t rotary_dim = arg_parser.get_int("rotary_dim"); bool i_perm = arg_parser.get_bool("iperm"); bool o_perm = arg_parser.get_bool("operm"); @@ -189,10 +174,7 @@ auto run(const ck_tile::ArgParser& arg_parser) hdim_q, hdim_v, seqlen_knew, - seqlen_qpads, seqlen_kpads, - q_eff_lens_per_batch, - kv_eff_lens_per_batch, rotary_dim, i_perm, o_perm, diff --git a/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp b/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp index 7ddb65a2db..569c98a458 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd_v3.cpp @@ -52,16 +52,7 @@ auto parse_cmd_args(int argc, char* argv[]) -> std::pair get_query_shape() const @@ -183,8 +172,6 @@ struct Problem mask_info mask; TensorLayout input_layout; TensorLayout output_layout; - std::vector q_eff_lens; - std::vector kv_eff_lens; }; struct RunConfig @@ -339,10 +326,8 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) q_buf.ToDevice(q.data()); k_buf.ToDevice(k.data()); v_buf.ToDevice(v.data()); - // Ensure output buffer is zero-initialized so padded regions compare cleanly - o_buf.SetZero(); - ck_tile::fmha_fwd_v3_args args{}; + ck_tile::fmha_fwd_v3_args args; args.data_type = problem.data_type; args.batch = problem.batch; @@ -395,60 +380,6 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) : problem.seqlen_q * problem.hdim; args.batch_stride_o = problem.seqlen_q * problem.nhead_q * problem.hdim; - // Optional cumulative seqlen overrides (exclude PAD) - const bool has_varlen_q = !problem.q_eff_lens.empty() && problem.q_eff_lens[0] != -1; - const bool has_varlen_k = !problem.kv_eff_lens.empty() && problem.kv_eff_lens[0] != -1; - - auto make_effective_vec = [&](const std::vector& opt_vec, ck_tile::index_t fallback) { - std::vector eff; - if(!opt_vec.empty() && opt_vec[0] != -1) - { - eff.assign(opt_vec.begin(), opt_vec.end()); - if(eff.size() < static_cast(problem.batch)) - { - eff.resize(problem.batch, eff.back()); - } - } - else - { - eff.assign(problem.batch, fallback); - } - return eff; - }; - - const auto eff_q_vec = make_effective_vec(problem.q_eff_lens, problem.seqlen_q); - const auto eff_kv_vec = make_effective_vec(problem.kv_eff_lens, problem.seqlen_k); - - // Calculate cumulative sums for kernel arguments if varlen is used - std::vector cuq_cum, cukv_cum; - auto calculate_cumulative = [&](const std::vector& per_batch_vec, - std::vector& cum_vec) { - cum_vec.resize(per_batch_vec.size() + 1); - cum_vec[0] = 0; - for(std::size_t i = 0; i < per_batch_vec.size(); ++i) - cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i]; - }; - - if(has_varlen_q) - { - calculate_cumulative(eff_q_vec, cuq_cum); - } - if(has_varlen_k) - { - calculate_cumulative(eff_kv_vec, cukv_cum); - } - - ck_tile::DeviceMem cuq_buf(!cuq_cum.empty() ? cuq_cum.size() * sizeof(ck_tile::index_t) : 0); - ck_tile::DeviceMem cukv_buf(!cukv_cum.empty() ? cukv_cum.size() * sizeof(ck_tile::index_t) : 0); - cuq_buf.ToDevice(!cuq_cum.empty() ? cuq_cum.data() : nullptr); - cukv_buf.ToDevice(!cukv_cum.empty() ? cukv_cum.data() : nullptr); - args.cu_seqlen_q_ptr = - !cuq_cum.empty() ? reinterpret_cast(cuq_buf.GetDeviceBuffer()) - : nullptr; - args.cu_seqlen_kv_ptr = - !cukv_cum.empty() ? reinterpret_cast(cukv_buf.GetDeviceBuffer()) - : nullptr; - ck_tile::stream_config stream_config{nullptr, true, /*log_level=*/0, @@ -511,72 +442,15 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) o_ref = o_ref.transpose({0, 2, 1, 3}); } - // If variable lengths are provided, compute per-batch references - // with the effective lengths; else compute a single full reference. - if(has_varlen_q || has_varlen_k) - { - // Variable-length aware verification: zero-fill padded region and only compute valid part. - o_ref.SetZero(); - - for(int b = 0; b < problem.batch; ++b) - { - const ck_tile::index_t seqlen_q_eff = eff_q_vec[b]; - const ck_tile::index_t seqlen_kv_eff = eff_kv_vec[b]; - - if(seqlen_q_eff <= 0 || seqlen_kv_eff <= 0) - continue; - - // Slice current batch from inputs (bshd) and build single-batch tensors - ck_tile::HostTensor q_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim}); - ck_tile::HostTensor k_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); - ck_tile::HostTensor v_b({1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); - ck_tile::HostTensor o_b({1, seqlen_q_eff, problem.nhead_q, problem.hdim}); - - // Copy effective region - q_b.ForEach([&](auto& self, auto idx) { - // idx: [0, s, h, d] - self(idx) = q(b, idx[1], idx[2], idx[3]); - }); - k_b.ForEach([&](auto& self, auto idx) { self(idx) = k(b, idx[1], idx[2], idx[3]); }); - v_b.ForEach([&](auto& self, auto idx) { self(idx) = v(b, idx[1], idx[2], idx[3]); }); - - // Compute reference for this batch segment (host::fmha_fwd expects bshd tensors) - host::fmha_fwd(q_b, - k_b, - v_b, - problem.mask, - o_b, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::scales{problem.softmax_scale}); - - // Scatter into o_ref's bshd descriptor memory - for(int s = 0; s < seqlen_q_eff; ++s) - { - for(int h = 0; h < problem.nhead_q; ++h) - { - for(int d = 0; d < problem.hdim; ++d) - { - o_ref(b, s, h, d) = o_b(0, s, h, d); - } - } - } - } - } - else - { - // No varlen override: compute the full reference once - host::fmha_fwd(q, - k, - v, - problem.mask, - o_ref, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::scales{problem.softmax_scale}); - } + host::fmha_fwd(q, + k, + v, + problem.mask, + o_ref, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::scales{problem.softmax_scale}); ck_tile::HostTensor o(problem.get_output_shape()); o_buf.FromDevice(o.data()); diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index f5dd42a6bd..c41e48e6aa 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -162,20 +162,11 @@ struct fmha_fwd_args void* lse_ptr; void* o_ptr; - // Optional cumulative sequence length arrays - // Batch mode: cu_seqlen_* override effective per-batch lengths (exclude PAD) - const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] - const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] - const void* seqstart_q_ptr; const void* seqstart_k_ptr; const void* seqlen_k_ptr; // only used if both 'seqstart_q_ptr' & 'seqstart_k_ptr' are not nullptr - // Group mode: seqstart_padded_* provide physical starts including PAD (optional) - const void* seqstart_padded_q_ptr = nullptr; // [batch+1] - const void* seqstart_padded_k_ptr = nullptr; // [batch+1] - ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; ck_tile::index_t batch; @@ -563,9 +554,7 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.min_seqlen_q, args.p_drop, args.s_randval, - args.drop_seed_offset, - args.seqstart_padded_q_ptr, - args.seqstart_padded_k_ptr); + args.drop_seed_offset); } else { // create batch mode kernel arguments @@ -611,9 +600,7 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset, - args.cu_seqlen_q_ptr, - args.cu_seqlen_kv_ptr); + args.drop_seed_offset); } }(); diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index cb5827975e..43f484fe14 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -151,10 +151,7 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t seqlen_knew, - std::vector seqlen_qpads, std::vector seqlen_kpads, - std::vector q_eff_lens_per_batch, - std::vector kv_eff_lens_per_batch, ck_tile::index_t rotary_dim, bool i_perm, bool o_perm, @@ -365,44 +362,6 @@ fwd_result fmha_fwd_run(mode_enum mode, const auto seqstart_k_host = to_seqstarts(seqlen_ks); const auto seqstart_k_with_padding_host = to_seqstarts(seqlen_kpads); - // Optional padded Q seqstarts (group-mode only) - std::vector seqstart_q_with_padding_host; - if(mode == mode_enum::group && !seqlen_qpads.empty() && seqlen_qpads[0] != -1) - { - if(seqlen_qpads.size() < static_cast(batch)) - { - seqlen_qpads.resize(batch, seqlen_qpads.back()); - } - if(seqlen_qpads.size() == static_cast(batch)) - { - seqstart_q_with_padding_host = to_seqstarts( - ck_tile::span(seqlen_qpads.data(), seqlen_qpads.size())); - } - } - - // Optional batch-mode cumulative seqlen overrides - std::vector cuq_cum, cukv_cum; - if(mode == mode_enum::batch) - { - auto calculate_cumulative = [&](std::vector& per_batch_vec, - std::vector& cum_vec) { - if(!per_batch_vec.empty() && per_batch_vec[0] != -1) - { - if(per_batch_vec.size() < static_cast(batch)) - { - per_batch_vec.resize(batch, per_batch_vec.back()); - } - cum_vec.resize(batch + 1); - cum_vec[0] = 0; - for(int i = 0; i < batch; ++i) - cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i]; - } - }; - - calculate_cumulative(q_eff_lens_per_batch, cuq_cum); - calculate_cumulative(kv_eff_lens_per_batch, cukv_cum); - } - using TypeConfig = FmhaFwdTypeConfig; using QDataType = typename TypeConfig::QDataType; @@ -486,15 +445,8 @@ fwd_result fmha_fwd_run(mode_enum mode, // host memory for storing all the tensor elements const ck_tile::index_t shape_batch = (mode == mode_enum::batch ? batch : 1); - // logical(unpadded) total seqlen_q for group; batch uses fixed seqlen - const ck_tile::index_t shape_seqlen_q_lse = - (mode == mode_enum::batch ? seqlen_qs[0] : seqstart_q_host.back()); - // physical(padded) total seqlen_q for group when s_qpad is provided; else use logical const ck_tile::index_t shape_seqlen_q = - (mode == mode_enum::batch - ? seqlen_qs[0] - : (seqstart_q_with_padding_host.empty() ? seqstart_q_host.back() - : seqstart_q_with_padding_host.back())); + (mode == mode_enum::batch ? seqlen_qs[0] : seqstart_q_host.back()); const ck_tile::index_t shape_seqlen_k = (mode == mode_enum::batch ? seqlen_ks[0] : (seqlen_kpads[0] < 0 ? seqstart_k_host.back() @@ -552,7 +504,7 @@ fwd_result fmha_fwd_run(mode_enum mode, // batch mode of lse data layout is [batch, nhead, seqlen_q] // group mode of lse data layout is [nhead, total_seqlen_q] ck_tile::HostTensor lse_host( - lse ? std::array{shape_batch, nhead, shape_seqlen_q_lse} + lse ? std::array{shape_batch, nhead, shape_seqlen_q} : std::array{1, 1, 1} /* dummy shape for simplifying code */); ck_tile::HostTensor o_host( @@ -650,16 +602,6 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t)); ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t)); - ck_tile::DeviceMem seqstart_q_padded_buf(seqstart_q_with_padding_host.empty() - ? 0 - : seqstart_q_with_padding_host.size() * - sizeof(int32_t)); - ck_tile::DeviceMem seqstart_k_padded_buf( - seqlen_kpads[0] < 0 ? 0 : seqstart_k_with_padding_host.size() * sizeof(int32_t)); - ck_tile::DeviceMem cu_seqlen_q_buf(cuq_cum.empty() ? 0 - : cuq_cum.size() * sizeof(ck_tile::index_t)); - ck_tile::DeviceMem cu_seqlen_kv_buf( - cukv_cum.empty() ? 0 : cukv_cum.size() * sizeof(ck_tile::index_t)); ck_tile::DeviceMem seqlen_k_buf((mode == mode_enum::batch && use_kvcache) || 0 <= seqlen_kpads[0] ? seqlen_ks.size() * sizeof(int32_t) @@ -751,14 +693,8 @@ fwd_result fmha_fwd_run(mode_enum mode, vnew_buf.ToDevice(vnew_host.data()); bias_buf.ToDevice(bias_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); - // Keep logical starts in seqstart_k; pass padded K via separate pointer - seqstart_k.ToDevice(seqstart_k_host.data()); - seqstart_q_padded_buf.ToDevice( - seqstart_q_with_padding_host.empty() ? nullptr : seqstart_q_with_padding_host.data()); - seqstart_k_padded_buf.ToDevice(seqlen_kpads[0] < 0 ? nullptr - : seqstart_k_with_padding_host.data()); - cu_seqlen_q_buf.ToDevice(cuq_cum.empty() ? nullptr : cuq_cum.data()); - cu_seqlen_kv_buf.ToDevice(cukv_cum.empty() ? nullptr : cukv_cum.data()); + seqstart_k.ToDevice(seqlen_kpads[0] < 0 ? seqstart_k_host.data() + : seqstart_k_with_padding_host.data()); seqlen_k_buf.ToDevice((mode == mode_enum::batch && use_kvcache) || 0 <= seqlen_kpads[0] ? seqlen_ks.data() : nullptr); @@ -894,8 +830,8 @@ fwd_result fmha_fwd_run(mode_enum mode, const ck_tile::index_t nhead_stride_bias = (i_perm ? 0 * shape_seqlen_q * max_seqlen_k : 0 * max_seqlen_k); const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t nhead_stride_lse = shape_seqlen_q_lse; - const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q_lse); + const ck_tile::index_t nhead_stride_lse = shape_seqlen_q; + const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q); const ck_tile::index_t nhead_stride_o_acc = (num_splits * shape_seqlen_q * hdim_v); const ck_tile::index_t nhead_stride_o = (o_perm ? shape_seqlen_q * hdim_v : hdim_v); // setup batch_stride_* arguments @@ -910,8 +846,8 @@ fwd_result fmha_fwd_run(mode_enum mode, const ck_tile::index_t batch_stride_vnew = (nhead_k * hdim_v * seqlen_knew); const ck_tile::index_t batch_stride_bias = (0 * nhead * shape_seqlen_q * max_seqlen_k); const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q_lse); - const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q_lse); + const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q); + const ck_tile::index_t batch_stride_lse_acc = (nhead * num_splits * shape_seqlen_q); const ck_tile::index_t batch_stride_o_acc = (nhead * num_splits * shape_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_o = (nhead * shape_seqlen_q * hdim_v); const ck_tile::index_t batch_stride_block_table = (max_num_page_blocks / batch); @@ -1025,29 +961,6 @@ fwd_result fmha_fwd_run(mode_enum mode, { args.drop_seed_offset = std::make_pair(drop_seed, drop_offset); } - - // Group-mode: optional physical padded starts for Q/K - if(mode == mode_enum::group) - { - args.seqstart_padded_q_ptr = (seqstart_q_with_padding_host.empty() - ? nullptr - : seqstart_q_padded_buf.GetDeviceBuffer()); - args.seqstart_padded_k_ptr = - (seqlen_kpads[0] < 0 ? nullptr : seqstart_k_padded_buf.GetDeviceBuffer()); - } - - // Batch-mode: optional cumulative effective seqlen overrides - if(mode == mode_enum::batch) - { - args.cu_seqlen_q_ptr = cuq_cum.empty() - ? nullptr - : reinterpret_cast( - cu_seqlen_q_buf.GetDeviceBuffer()); - args.cu_seqlen_kv_ptr = cukv_cum.empty() - ? nullptr - : reinterpret_cast( - cu_seqlen_kv_buf.GetDeviceBuffer()); - } } else if constexpr(std::is_same_v>) { @@ -1254,29 +1167,15 @@ fwd_result fmha_fwd_run(mode_enum mode, for(ck_tile::index_t wb = 0; wb < batch; ++wb) { - ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb]; - ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb]; - if(mode == mode_enum::batch) - { - if(!cuq_cum.empty()) - { - real_seqlen_q = cuq_cum[wb + 1] - cuq_cum[wb]; - } - if(!cukv_cum.empty()) - { - real_seqlen_k = cukv_cum[wb + 1] - cukv_cum[wb]; - } - } + const ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb]; + const ck_tile::index_t real_seqlen_k = seqstart_k_host[wb + 1] - seqstart_k_host[wb]; // adjust matrix index according to the mode const ck_tile::index_t b_idx = (mode == mode_enum::batch ? wb : 0); const ck_tile::index_t cache_b_idx = (use_cache_batch_idx ? cache_batch_idx_host(b_idx) : b_idx); const ck_tile::index_t query_offset = - (mode == mode_enum::batch - ? 0 - : (seqstart_q_with_padding_host.empty() ? seqstart_q_host[wb] - : seqstart_q_with_padding_host[wb])); + (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]); const ck_tile::index_t key_offset = (mode == mode_enum::batch ? 0 @@ -1639,10 +1538,8 @@ fwd_result fmha_fwd_run(mode_enum mode, if(lse) { ck_tile::HostTensor lse_host_result({nhead, real_seqlen_q}); - const ck_tile::index_t query_offset_lse = - (mode == mode_enum::batch ? 0 : seqstart_q_host[wb]); lse_host_result.ForEach([&](auto& self, auto idx) { - self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset_lse); + self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset); }); cur_pass = ck_tile::check_err(lse_host_result, diff --git a/example/ck_tile/01_fmha/fmha_fwd_v3.hpp b/example/ck_tile/01_fmha/fmha_fwd_v3.hpp index 4bd1d1a367..10cb5149a4 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_v3.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_v3.hpp @@ -56,11 +56,6 @@ struct fmha_fwd_v3_args index_t stride_o; index_t nhead_stride_o; index_t batch_stride_o; - - // Optional batch-mode cumulative seqlen overrides (exclude PAD) - // If provided, they override per-batch effective lengths to skip tail padding. - const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] - const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] }; std::ostream& operator<<(std::ostream& stream, const fmha_fwd_v3_args::data_type_enum& data_type); diff --git a/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp b/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp index 194675f962..e0fbad39a5 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_v3_impl.hpp @@ -158,9 +158,7 @@ float fmha_fwd_v3_kernel_launch(const fmha_fwd_v3_args& args, const stream_confi args.window_size_left, args.window_size_right, args.mask_type, - remap_opt, - args.cu_seqlen_q_ptr, - args.cu_seqlen_kv_ptr); + remap_opt); dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, args.hdim_v); constexpr dim3 blocks = Kernel::BlockSize(); diff --git a/example/ck_tile/01_fmha/script/benchmark_fwd.sh b/example/ck_tile/01_fmha/script/benchmark_fwd.sh index 31ad800039..88c16cceb6 100755 --- a/example/ck_tile/01_fmha/script/benchmark_fwd.sh +++ b/example/ck_tile/01_fmha/script/benchmark_fwd.sh @@ -18,36 +18,3 @@ $EXE -prec=$prec -b=1 -h=$nhead -d=$hdim -s=16384 -iperm=$perm -operm=$perm -kn done done done - -#Padding Benchmarks: batch mode (baseline vs low/med/high pad) -prec="fp16" -base_batch_args="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID" - -# baseline (no pad) -$EXE $base_batch_args - -# low pad (≈90–95% effective) -$EXE $base_batch_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 - -# medium pad (≈60–75% effective) -$EXE $base_batch_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 - -# high pad (≈30–40% effective) -$EXE $base_batch_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 - -# Padding Benchmarks: group mode (baseline vs low/med/high physical pad) -seqlens_q="1024,768,512,256" -seqlens_k="1024,768,512,256" -base_group_args="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=1 -v=$VALID" - -# baseline (no physical pad) -$EXE $base_group_args - -# low physical pad -$EXE $base_group_args -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320 - -# medium physical pad -$EXE $base_group_args -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384 - -# high physical pad -$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 diff --git a/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh b/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh index a3f7d68eb3..b847e85398 100755 --- a/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh +++ b/example/ck_tile/01_fmha/script/benchmark_fwd_v3.sh @@ -23,20 +23,3 @@ done done done done - -# Padding benchmark comparisons for v3 (batch mode only) -# ==== V3 Padding Benchmarks: batch mode (baseline vs low/med/high pad) ==== -prec="fp16" -base_v3_args="-prec=$prec -b=4 -h=16 -d=128 -s=1024 -mask=0 -iperm=0 -operm=0 -v=$VALID" - -# baseline (no pad) -$EXE $base_v3_args - -# low pad (≈90–95% effective) -$EXE $base_v3_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 - -# medium pad (≈60–75% effective) -$EXE $base_v3_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 - -# high pad (≈30–40% effective) -$EXE $base_v3_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 diff --git a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh index fca6b8d0cd..afd0c728c6 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh @@ -137,118 +137,9 @@ run_fp16_appendkv_tests() { done ; done ; done } -run_padding_smoke_tests() { - # Padding-only smoke tests for batch/group mode using COMMON_ARGS - local prec="fp16" - - # Batch mode: padding via effective lengths (exclude PAD) - # Use lse=1 to select a non-trload kernel and avoid overly strict tolerance mismatches - local base_batch="-prec=$prec -mode=0 -b=4 -h=16 -h_k=16 -d=128 -s=1024 -bias=n -mask=0 -lse=1 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS" - # low pad (≈90–95% effective) - $EXE $base_batch -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 - # medium pad (≈60–75% effective) - $EXE $base_batch -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 - # high pad (≈30–40% effective) - $EXE $base_batch -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 - - # Group mode: padding via physical stride along seqlen - local seqlens_q="1024,768,512,256" - local seqlens_k="1024,768,512,256" - local base_group="-prec=$prec -mode=1 -b=4 -h=16 -h_k=16 -d=128 -s=$seqlens_q -s_k=$seqlens_k -bias=n -mask=0 -lse=0 -iperm=0 -operm=0 -vlayout=r -kname=$KNAME $COMMON_ARGS" - # low physical pad - $EXE $base_group -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320 - # medium physical pad - $EXE $base_group -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384 - # high physical pad - $EXE $base_group -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512 -} - -run_padding_basic_boundary_tests() { - # Basic padding and boundary tests (reference: smoke_test_fwd_pad.sh) - local prec - local perm - - # Group mode: Q&K padded with per-batch different strides - for prec in fp16 bf16 ; do - for perm in 0 1 ; do - $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=16 -d_v=32 \ - -s=55 -s_k=256 -s_qpad=64,60 -s_kpad=272,260 \ - -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \ - -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS - done - done - - # slightly larger, uneven padding strides - for prec in fp16 bf16 ; do - for perm in 0 1 ; do - $EXE -prec=$prec -mode=1 -b=3 -h=2 -h_k=1 -d=64 -d_v=64 \ - -s=50,60,40 -s_k=128,256,192 -s_qpad=64,64,64 -s_kpad=160,288,224 \ - -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ - -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS - done - done - - # only K padded; Q unpadded - for prec in fp16 bf16 ; do - for perm in 0 1 ; do - $EXE -prec=$prec -mode=1 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 \ - -s=55 -s_k=256 -s_kpad=272,260 \ - -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ - -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS - done - done - - # use cu_seqlen overrides to skip tail PAD - for prec in fp16 bf16 ; do - for perm in 0 1 ; do - $EXE -prec=$prec -mode=0 -b=4 -h=8 -h_k=8 -d=128 -s=3 -s_k=3 \ - -q_eff_lens=1,2,1,2 -kv_eff_lens=1,2,1,2 \ - -bias=n -p_drop=0.0 -lse=1 -iperm=$perm -operm=$perm \ - -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS - - $EXE -prec=$prec -mode=0 -b=2 -h=2 -h_k=1 -d=32 -d_v=64 -s=64 -s_k=256 \ - -q_eff_lens=55,60 -kv_eff_lens=200,256 \ - -bias=n -p_drop=0.0 -lse=0 -iperm=$perm -operm=$perm \ - -num_splits=1 -page_block_size=0 -cache_batch_idx=0 -kname=$KNAME $COMMON_ARGS - done - done - - # no padding (equal), mixed Q/KV, all len=1 - for prec in fp16 bf16 ; do - $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ - -q_eff_lens=128,128,128,128 -kv_eff_lens=128,128,128,128 \ - -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS - - $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ - -q_eff_lens=10,20,30,40 -kv_eff_lens=40,30,20,10 \ - -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS - - $EXE -prec=$prec -mode=0 -b=4 -h=8 -d=64 -s=128 -s_k=128 \ - -q_eff_lens=1,1,1,1 -kv_eff_lens=1,1,1,1 \ - -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS - done - - # highly variable logical lengths - for prec in fp16 bf16 ; do - $EXE -prec=$prec -mode=1 -b=4 -h=4 -d=32 \ - -s=1,127,3,65 -s_k=1,127,3,65 -s_kpad=128 \ - -bias=n -p_drop=0.0 -lse=1 -kname=$KNAME $COMMON_ARGS - done - - # GQA + Alibi + Causal mask (keep vlayout row-major for fp16/bf16 - for prec in fp16 bf16 ; do - $EXE -prec=$prec -mode=1 -b=2 -h=16 -h_k=4 -d=128 \ - -s=256,129 -s_k=256,129 -s_kpad=256 \ - -bias=a -mask=t -lse=1 -iperm=0 -operm=0 -vlayout=r \ - -kname=$KNAME $COMMON_ARGS - done -} - set -x run_fp16_bf16_tests -run_padding_smoke_tests -run_padding_basic_boundary_tests run_fp8_tests run_fp8bf16_tests run_fp8fp32_tests diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 3f417bc125..58fdad149a 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -291,11 +291,6 @@ struct FmhaFwdKernel ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; ck_tile::index_t batch_stride_o; - - // Optional cumulative sequence length pointers for batch mode - // If provided, they override seqlen_q / seqlen_k per-batch to skip tail padding. - const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // cumulative, length without PAD - const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // cumulative, length without PAD }; struct FmhaFwdGroupModeKargs @@ -315,11 +310,6 @@ struct FmhaFwdKernel const int32_t* seqstart_q_ptr; const int32_t* seqstart_k_ptr; const int32_t* seqlen_k_ptr; - - // Optional cumulative padded sequence starts (including PAD tokens) - // Used solely to compute memory offsets when sequences are physically padded. - const int32_t* seqstart_padded_q_ptr = nullptr; - const int32_t* seqstart_padded_k_ptr = nullptr; }; using Kargs = std::conditional_t; @@ -470,105 +460,6 @@ struct FmhaFwdKernel return kargs; } - // Overload: Batch mode with optional cu_seqlen pointers (unpadded cumulative lengths) - template - CK_TILE_HOST static constexpr std::enable_if_t - MakeKargsImpl(const void* q_ptr, - const void* k_ptr, - const void* v_ptr, - const void* bias_ptr, - void* rand_val_ptr, - void* lse_ptr, - void* o_ptr, - ck_tile::index_t seqlen_q, - ck_tile::index_t seqlen_k, - ck_tile::index_t hdim_q, - ck_tile::index_t hdim_v, - ck_tile::index_t num_head_q, - ck_tile::index_t nhead_ratio_qk, - float scale_s, - float scale_p, - float scale_o, - float logits_soft_cap, - ck_tile::index_t stride_q, - ck_tile::index_t stride_k, - ck_tile::index_t stride_v, - ck_tile::index_t stride_bias, - ck_tile::index_t stride_randval, - ck_tile::index_t stride_o, - ck_tile::index_t nhead_stride_q, - ck_tile::index_t nhead_stride_k, - ck_tile::index_t nhead_stride_v, - ck_tile::index_t nhead_stride_bias, - ck_tile::index_t nhead_stride_randval, - ck_tile::index_t nhead_stride_lse, - ck_tile::index_t nhead_stride_o, - ck_tile::index_t batch_stride_q, - ck_tile::index_t batch_stride_k, - ck_tile::index_t batch_stride_v, - ck_tile::index_t batch_stride_bias, - ck_tile::index_t batch_stride_randval, - ck_tile::index_t batch_stride_lse, - ck_tile::index_t batch_stride_o, - ck_tile::index_t window_size_left, - ck_tile::index_t window_size_right, - ck_tile::index_t mask_type, - float p_drop, - bool s_randval, - std::variant, std::pair> - drop_seed_offset, - const ck_tile::index_t* cu_seqlen_q_ptr, - const ck_tile::index_t* cu_seqlen_kv_ptr) - { - auto kargs = MakeKargsImpl(q_ptr, - k_ptr, - v_ptr, - bias_ptr, - rand_val_ptr, - lse_ptr, - o_ptr, - seqlen_q, - seqlen_k, - hdim_q, - hdim_v, - num_head_q, - nhead_ratio_qk, - scale_s, - scale_p, - scale_o, - logits_soft_cap, - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_bias, - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - batch_stride_bias, - batch_stride_randval, - batch_stride_lse, - batch_stride_o, - window_size_left, - window_size_right, - mask_type, - p_drop, - s_randval, - drop_seed_offset); - - kargs.cu_seqlen_q_ptr = cu_seqlen_q_ptr; - kargs.cu_seqlen_kv_ptr = cu_seqlen_kv_ptr; - return kargs; - } - // std::variant<> can't take in a list initializer, overload for backward compatibility template CK_TILE_HOST static constexpr std::enable_if_t @@ -890,95 +781,6 @@ struct FmhaFwdKernel return kargs; } - // Overload: Group mode with optional padded seqstarts for memory offsets - template - CK_TILE_HOST static constexpr std::enable_if_t - MakeKargsImpl(const void* q_ptr, - const void* k_ptr, - const void* v_ptr, - const void* bias_ptr, - void* rand_val_ptr, - void* lse_ptr, - void* o_ptr, - const void* seqstart_q_ptr, - const void* seqstart_k_ptr, - const void* seqlen_k_ptr, - ck_tile::index_t hdim_q, - ck_tile::index_t hdim_v, - ck_tile::index_t num_head_q, - ck_tile::index_t nhead_ratio_qk, - float scale_s, - float scale_p, - float scale_o, - float logits_soft_cap, - ck_tile::index_t stride_q, - ck_tile::index_t stride_k, - ck_tile::index_t stride_v, - ck_tile::index_t stride_bias, - ck_tile::index_t stride_randval, - ck_tile::index_t stride_o, - ck_tile::index_t nhead_stride_q, - ck_tile::index_t nhead_stride_k, - ck_tile::index_t nhead_stride_v, - ck_tile::index_t nhead_stride_bias, - ck_tile::index_t nhead_stride_randval, - ck_tile::index_t nhead_stride_lse, - ck_tile::index_t nhead_stride_o, - ck_tile::index_t window_size_left, - ck_tile::index_t window_size_right, - ck_tile::index_t mask_type, - ck_tile::index_t min_seqlen_q, - float p_drop, - bool s_randval, - std::variant, std::pair> - drop_seed_offset, - const void* seqstart_padded_q_ptr, - const void* seqstart_padded_k_ptr) - { - auto kargs = MakeKargsImpl(q_ptr, - k_ptr, - v_ptr, - bias_ptr, - rand_val_ptr, - lse_ptr, - o_ptr, - seqstart_q_ptr, - seqstart_k_ptr, - seqlen_k_ptr, - hdim_q, - hdim_v, - num_head_q, - nhead_ratio_qk, - scale_s, - scale_p, - scale_o, - logits_soft_cap, - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_bias, - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - window_size_left, - window_size_right, - mask_type, - min_seqlen_q, - p_drop, - s_randval, - drop_seed_offset); - - kargs.seqstart_padded_q_ptr = reinterpret_cast(seqstart_padded_q_ptr); - kargs.seqstart_padded_k_ptr = reinterpret_cast(seqstart_padded_k_ptr); - return kargs; - } - // std::variant<> can't take in a list initializer, overload for backward compatibility template CK_TILE_HOST static constexpr std::enable_if_t @@ -1271,44 +1073,35 @@ struct FmhaFwdKernel if constexpr(kIsGroupMode) { - // logical and physical (padded) starts - const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; + // get starting offset for each batch + const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; - const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr - ? kargs.seqstart_padded_q_ptr[i_batch] - : query_start_unpadded; - const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr - ? kargs.seqstart_padded_k_ptr[i_batch] - : key_start_unpadded; - - // DRAM base offsets use physical padded starts - batch_offset_q = query_start_padded * kargs.stride_q; - batch_offset_k = key_start_padded * kargs.stride_k; + batch_offset_q = query_start * kargs.stride_q; + batch_offset_k = key_start * kargs.stride_k; if constexpr(std::is_same_v) { - batch_offset_v = key_start_padded * kargs.stride_v; + batch_offset_v = key_start * kargs.stride_v; } else { - batch_offset_v = key_start_padded; + batch_offset_v = key_start; } if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { - batch_offset_bias = query_start_padded * kargs.stride_bias; + batch_offset_bias = query_start * kargs.stride_bias; } if constexpr(kStoreLSE) { - // LSE stays indexed by unpadded starts - batch_offset_lse = query_start_unpadded; + batch_offset_lse = query_start; } if constexpr(kHasDropout) { - batch_offset_randval = query_start_padded * kargs.stride_randval; + batch_offset_randval = query_start * kargs.stride_randval; } - batch_offset_o = query_start_padded * kargs.stride_o; + batch_offset_o = query_start * kargs.stride_o; - // real logical lengths (exclude PAD) + // get real # queries & # keys under group mode const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0]; @@ -1320,7 +1113,8 @@ struct FmhaFwdKernel } } - // terminate unnecessary blocks earlier + // # of required blocks is different in each groups, terminate unnecessary blocks + // earlier if(kargs.seqlen_q <= i_m0) { return; @@ -1356,18 +1150,6 @@ struct FmhaFwdKernel static_cast(i_batch) * kargs.batch_stride_randval; } batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; - - // If cumulative seqlen pointers are provided, override per-batch effective lengths - if(kargs.cu_seqlen_q_ptr != nullptr) - { - kargs.seqlen_q = - kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; - } - if(kargs.cu_seqlen_kv_ptr != nullptr) - { - kargs.seqlen_k = - kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; - } } // for simplicity, batch stride we just modify the pointer @@ -1766,35 +1548,26 @@ struct FmhaFwdKernel if constexpr(kIsGroupMode) { // get starting offset for each batch - const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; + const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; - const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr - ? kargs.seqstart_padded_q_ptr[i_batch] - : query_start_unpadded; - const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr - ? kargs.seqstart_padded_k_ptr[i_batch] - : key_start_unpadded; - - batch_offset_q = query_start_padded * kargs.stride_q; - batch_offset_k = key_start_padded * kargs.stride_k; + batch_offset_q = query_start * kargs.stride_q; + batch_offset_k = key_start * kargs.stride_k; if constexpr(std::is_same_v) { - batch_offset_v = key_start_padded * kargs.stride_v; + batch_offset_v = key_start * kargs.stride_v; } else { - // col-major V: offset along seqlen dimension is scalar index - batch_offset_v = key_start_padded; + batch_offset_v = key_start; } if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { - batch_offset_bias = query_start_padded * kargs.stride_bias; + batch_offset_bias = query_start * kargs.stride_bias; } - // LSE layout is [nhead, total_seqlen], index by unpadded start - batch_offset_lse = query_start_unpadded; - batch_offset_o = query_start_padded * kargs.stride_o; + batch_offset_lse = query_start; + batch_offset_o = query_start * kargs.stride_o; // get real # queries & # keys under group mode kargs.seqlen_q = kargs.seqstart_q_ptr[i_batch + 1] - kargs.seqstart_q_ptr[i_batch]; @@ -1832,18 +1605,6 @@ struct FmhaFwdKernel batch_offset_bias = static_cast(i_batch) * kargs.batch_stride_bias; } - - // If cumulative seqlen pointers are provided, override per-batch effective lengths - if(kargs.cu_seqlen_q_ptr != nullptr) - { - kargs.seqlen_q = - kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; - } - if(kargs.cu_seqlen_kv_ptr != nullptr) - { - kargs.seqlen_k = - kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; - } } // for simplicity, batch stride we just modify the pointer diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp index 52b9da40b8..c5e5745817 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp @@ -100,11 +100,6 @@ struct FmhaFwdV3Kernel ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; ck_tile::index_t batch_stride_o; - - // Optional cumulative sequence length pointers for batch mode - // If provided, they override seqlen_q / seqlen_k per-batch to skip tail padding. - const ck_tile::index_t* cu_seqlen_q_ptr = nullptr; // [batch+1] - const ck_tile::index_t* cu_seqlen_kv_ptr = nullptr; // [batch+1] }; struct FmhaFwdGroupModeKargs @@ -115,11 +110,6 @@ struct FmhaFwdV3Kernel const int32_t* seqstart_q_ptr; const int32_t* seqstart_k_ptr; const int32_t* seqlen_k_ptr; - - // Optional cumulative padded sequence starts (including PAD tokens) - // Used solely to compute memory offsets when sequences are physically padded. - const int32_t* seqstart_padded_q_ptr = nullptr; // [batch+1] - const int32_t* seqstart_padded_k_ptr = nullptr; // [batch+1] }; using Kargs = std::conditional_t; @@ -200,78 +190,6 @@ struct FmhaFwdV3Kernel return kargs; } - // Overload: Batch mode with optional cu_seqlen pointers - template - CK_TILE_HOST static constexpr std::enable_if_t - MakeKargs(const void* q_ptr, - const void* k_ptr, - const void* v_ptr, - void* lse_ptr, - void* o_ptr, - ck_tile::index_t seqlen_q, - ck_tile::index_t seqlen_k, - ck_tile::index_t hdim_q, - ck_tile::index_t hdim_v, - ck_tile::index_t num_head_q, - ck_tile::index_t nhead_ratio_qk, - float scale_s, - ck_tile::index_t stride_q, - ck_tile::index_t stride_k, - ck_tile::index_t stride_v, - ck_tile::index_t stride_o, - ck_tile::index_t nhead_stride_q, - ck_tile::index_t nhead_stride_k, - ck_tile::index_t nhead_stride_v, - ck_tile::index_t nhead_stride_lse, - ck_tile::index_t nhead_stride_o, - ck_tile::index_t batch_stride_q, - ck_tile::index_t batch_stride_k, - ck_tile::index_t batch_stride_v, - ck_tile::index_t batch_stride_lse, - ck_tile::index_t batch_stride_o, - ck_tile::index_t window_size_left, - ck_tile::index_t window_size_right, - ck_tile::index_t mask_type, - ck_tile::index_t remap_opt, - const ck_tile::index_t* cu_seqlen_q_ptr, - const ck_tile::index_t* cu_seqlen_kv_ptr) - { - auto kargs = MakeKargs(q_ptr, - k_ptr, - v_ptr, - lse_ptr, - o_ptr, - seqlen_q, - seqlen_k, - hdim_q, - hdim_v, - num_head_q, - nhead_ratio_qk, - scale_s, - stride_q, - stride_k, - stride_v, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - batch_stride_lse, - batch_stride_o, - window_size_left, - window_size_right, - mask_type, - remap_opt); - - kargs.cu_seqlen_q_ptr = cu_seqlen_q_ptr; - kargs.cu_seqlen_kv_ptr = cu_seqlen_kv_ptr; - return kargs; - } - template CK_TILE_HOST static constexpr std::enable_if_t MakeKargs(const void* q_ptr, @@ -342,70 +260,6 @@ struct FmhaFwdV3Kernel return kargs; } - // Overload: Group mode with optional padded seqstarts for memory offsets - template - CK_TILE_HOST static constexpr std::enable_if_t - MakeKargs(const void* q_ptr, - const void* k_ptr, - const void* v_ptr, - void* lse_ptr, - void* o_ptr, - const void* seqstart_q_ptr, - const void* seqstart_k_ptr, - const void* seqlen_k_ptr, - ck_tile::index_t hdim_q, - ck_tile::index_t hdim_v, - ck_tile::index_t num_head_q, - ck_tile::index_t nhead_ratio_qk, - float scale_s, - ck_tile::index_t stride_q, - ck_tile::index_t stride_k, - ck_tile::index_t stride_v, - ck_tile::index_t stride_o, - ck_tile::index_t nhead_stride_q, - ck_tile::index_t nhead_stride_k, - ck_tile::index_t nhead_stride_v, - ck_tile::index_t nhead_stride_lse, - ck_tile::index_t nhead_stride_o, - ck_tile::index_t window_size_left, - ck_tile::index_t window_size_right, - ck_tile::index_t mask_type, - ck_tile::index_t remap_opt, - const void* seqstart_padded_q_ptr, - const void* seqstart_padded_k_ptr) - { - auto kargs = MakeKargs(q_ptr, - k_ptr, - v_ptr, - lse_ptr, - o_ptr, - seqstart_q_ptr, - seqstart_k_ptr, - seqlen_k_ptr, - hdim_q, - hdim_v, - num_head_q, - nhead_ratio_qk, - scale_s, - stride_q, - stride_k, - stride_v, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_lse, - nhead_stride_o, - window_size_left, - window_size_right, - mask_type, - remap_opt); - - kargs.seqstart_padded_q_ptr = reinterpret_cast(seqstart_padded_q_ptr); - kargs.seqstart_padded_k_ptr = reinterpret_cast(seqstart_padded_k_ptr); - return kargs; - } - CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_, @@ -519,26 +373,18 @@ struct FmhaFwdV3Kernel if constexpr(kIsGroupMode) { // get starting offset for each batch - const long_index_t query_start_unpadded = kargs.seqstart_q_ptr[i_batch]; - const long_index_t key_start_unpadded = kargs.seqstart_k_ptr[i_batch]; + const long_index_t query_start = kargs.seqstart_q_ptr[i_batch]; + const long_index_t key_start = kargs.seqstart_k_ptr[i_batch]; - const long_index_t query_start_padded = kargs.seqstart_padded_q_ptr - ? kargs.seqstart_padded_q_ptr[i_batch] - : query_start_unpadded; - const long_index_t key_start_padded = kargs.seqstart_padded_k_ptr - ? kargs.seqstart_padded_k_ptr[i_batch] - : key_start_unpadded; - - batch_offset_q = query_start_padded * kargs.stride_q; - batch_offset_k = key_start_padded * kargs.stride_k; - batch_offset_v = key_start_padded * kargs.stride_v; + batch_offset_q = query_start * kargs.stride_q; + batch_offset_k = key_start * kargs.stride_k; + batch_offset_v = key_start * kargs.stride_v; if constexpr(kStoreLSE) { - // LSE layout is [nhead, total_seqlen], index by unpadded start - batch_offset_lse = query_start_unpadded; + batch_offset_lse = query_start; } - batch_offset_o = query_start_padded * kargs.stride_o; + batch_offset_o = query_start * kargs.stride_o; // get real # queries & # keys under group mode const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch; @@ -571,18 +417,6 @@ struct FmhaFwdV3Kernel batch_offset_lse = static_cast(i_batch) * kargs.batch_stride_lse; } batch_offset_o = static_cast(i_batch) * kargs.batch_stride_o; - - // If cumulative seqlen pointers are provided, override per-batch effective lengths - if(kargs.cu_seqlen_q_ptr != nullptr) - { - kargs.seqlen_q = - kargs.cu_seqlen_q_ptr[i_batch + 1] - kargs.cu_seqlen_q_ptr[i_batch]; - } - if(kargs.cu_seqlen_kv_ptr != nullptr) - { - kargs.seqlen_k = - kargs.cu_seqlen_kv_ptr[i_batch + 1] - kargs.cu_seqlen_kv_ptr[i_batch]; - } } // for simplicity, batch stride we just modify the pointer diff --git a/test/ck_tile/fmha/test_fmha_fwd.inc b/test/ck_tile/fmha/test_fmha_fwd.inc index 66d4e3dc21..08abd3358d 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.inc +++ b/test/ck_tile/fmha/test_fmha_fwd.inc @@ -98,10 +98,7 @@ TEST_P(AllLong, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim perm, // i_perm perm, // o_perm @@ -163,10 +160,7 @@ TEST_P(HDimPadding, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {seqlen_kpad}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim perm, // i_perm perm, // o_perm @@ -223,10 +217,7 @@ TEST_P(ElementwiseBias, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -282,10 +273,7 @@ TEST_P(Alibi, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim true, // i_perm true, // o_perm @@ -343,10 +331,7 @@ TEST_P(Dropout, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim false, // i_perm false, // o_perm @@ -406,10 +391,7 @@ TEST_P(PagedKV, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -475,10 +457,7 @@ TEST_P(SplitKV, Test) hdim_q, hdim_v, 0, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm false, // o_perm @@ -550,10 +529,7 @@ TEST_P(AppendKV, Test) hdim_q, hdim_v, seqlen_knew, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch 0, // rotary_dim i_perm, // i_perm true, // o_perm @@ -623,10 +599,7 @@ TEST_P(AppendKVRoPE, Test) hdim_q, hdim_v, seqlen_knew, // seqlen_knew - {-1}, // seqlen_qpads {-1}, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch rotary_dim, // rotary_dim i_perm, // i_perm true, // o_perm @@ -650,117 +623,3 @@ TEST_P(AppendKVRoPE, Test) } #endif // CK_TILE_FMHA_FWD_APPENDKV_API - -// --------------------------------------------------------------- -// Additional padding tests (q/kv physical padding & effective len) -// --------------------------------------------------------------- - -// Simple batch-mode test with per-batch Q/KV padding strides and effective lengths -TEST(TestCkTileFmhaFwd, BatchModeQKvPadding) -{ - if constexpr(std::is_same_v) - { - GTEST_SKIP() << "Skip for fp8"; - } - const mode_enum mode = mode_enum::batch; - const int batch = 3; - const int nhead = 2; - const int nhead_k = -1; - const int seqlen_q = 128; - const int seqlen_k = 128; - const int hdim_q = 64; - const int hdim_v = 64; - const int seqlen_knew = 0; - const std::vector seqlen_qpads{}; - const std::vector seqlen_kpads{}; - const std::vector q_eff_lens{120, 128, 100}; - const std::vector kv_eff_lens{110, 128, 90}; - - auto result = fmha_fwd_run(mode, - batch, - nhead, - nhead_k, - {adjust_seqlen(seqlen_q)}, - {adjust_seqlen(seqlen_k)}, - hdim_q, - hdim_v, - seqlen_knew, // seqlen_knew - seqlen_qpads, // seqlen_qpads - seqlen_kpads, // seqlen_kpads - q_eff_lens, // q_eff_lens_per_batch - kv_eff_lens, // kv_eff_lens_per_batch - 0, // rotary_dim - true, // i_perm - true, // o_perm - 0, // scale_s - 0, // logits_soft_cap - def_is_v_rowmajor, - def_lse, // lse - 0, // page_block_size - false, // use_cache_batch_idx - "n", // bias_str - 0.0f, // p_drop - 0, // drop_seed - 0, // drop_offset - false, // drop_prefs - "0", // mask_str - QUANT_ARGS, - true, // is_rotary_interleaved - 1, // num_splits - COMMON_ARGS); - CHECK_RESULT(result); -} - -// Simple group-mode test with uniform seqlen but per-batch padding & effective lengths -TEST(TestCkTileFmhaFwd, GroupModeQKvPadding) -{ - if constexpr(std::is_same_v) - { - GTEST_SKIP() << "Skip for fp8"; - } - const mode_enum mode = mode_enum::group; - const int batch = 2; - const int nhead = 2; - const int nhead_k = -1; - const std::vector seqlen_q{96, 128}; // unpadded - const std::vector seqlen_k{96, 128}; // unpadded - const int hdim_q = 64; - const int hdim_v = 64; - const int seqlen_knew = 0; - const std::vector seqlen_qpads{128, 160}; - const std::vector seqlen_kpads{128, 160}; - - auto result = fmha_fwd_run(mode, - batch, - nhead, - nhead_k, - seqlen_q, - seqlen_k, - hdim_q, - hdim_v, - seqlen_knew, // seqlen_knew - seqlen_qpads, // seqlen_qpads - seqlen_kpads, // seqlen_kpads - {}, // q_eff_lens_per_batch - {}, // kv_eff_lens_per_batch - 0, // rotary_dim - true, // i_perm - true, // o_perm - 0, // scale_s - 0, // logits_soft_cap - def_is_v_rowmajor, - def_lse, // lse - 0, // page_block_size - false, // use_cache_batch_idx - "n", // bias_str - 0.0f, // p_drop - 0, // drop_seed - 0, // drop_offset - false, // drop_prefs - "0", // mask_str - QUANT_ARGS, - true, // is_rotary_interleaved - 1, // num_splits - COMMON_ARGS); - CHECK_RESULT(result); -} From 4363a82bd658542c063b419896be8b8826b61985 Mon Sep 17 00:00:00 2001 From: Sami Remes Date: Sat, 20 Sep 2025 02:52:35 +0300 Subject: [PATCH 10/12] [CK_TILE] Tensor-wise scaled quant gemm kernel (#2846) * rename gemm_group_quant to gemm_quant * Add TensorWise quant mode * Cshuffle epilogue tests with tensor scaling * Add tensor quant to example * Don't use readfirstlane for reading scales - doesn't work for some reason * Add to changelog * revert include - from a merge problem? * revert common.hpp include * revert host.hpp include * remove unused utility function * rename quant pipeline problem * refactor quant tests * remove aquant utils * use TEST_F * fix all tests by changing gemm config * Use typed tests * fix copyright --- CHANGELOG.md | 1 + .../17_grouped_gemm/quant_grouped_gemm.cpp | 20 +- example/ck_tile/38_block_scale_gemm/README.md | 4 +- .../38_block_scale_gemm/gemm_quant_basic.cpp | 49 +- .../38_block_scale_gemm/gemm_utils.hpp | 4 +- .../run_gemm_quant_example.inc | 85 +- include/ck_tile/core/tensor/load_tile.hpp | 5 +- .../ck_tile/host/reference/reference_gemm.hpp | 56 +- .../ops/epilogue/cshuffle_epilogue.hpp | 89 +- include/ck_tile/ops/gemm_group_quant.hpp | 21 - include/ck_tile/ops/gemm_quant.hpp | 21 + .../block_universal_gemm_as_aquant_bs_cr.hpp | 0 .../block_universal_gemm_as_bs_bquant_cr.hpp | 0 .../kernel/gemm_quant_kernel.hpp | 19 +- .../kernel/grouped_gemm_quant_kernel.hpp | 2 +- .../gemm_aquant_pipeline_ag_bg_cr_base.hpp | 0 .../gemm_aquant_pipeline_ag_bg_cr_policy.hpp | 0 .../gemm_aquant_pipeline_ag_bg_cr_v3.hpp | 2 +- .../gemm_bquant_pipeline_ag_bg_cr_base.hpp | 0 .../gemm_bquant_pipeline_ag_bg_cr_policy.hpp | 0 .../gemm_bquant_pipeline_ag_bg_cr_v3.hpp | 2 +- .../pipeline/gemm_group_quant_utils.hpp | 0 .../pipeline/gemm_quant_pipeline_problem.hpp | 27 +- .../pipeline/tile_gemm_quant_traits.hpp | 15 +- .../epilogue/test_cshuffle_epilogue.cpp | 45 +- .../epilogue/test_cshuffle_epilogue_util.hpp | 48 +- test/ck_tile/gemm_block_scale/CMakeLists.txt | 11 +- .../test_gemm_aquant_basic_bf8.cpp | 6 - .../test_gemm_aquant_basic_fp8.cpp | 6 - .../test_gemm_aquant_basic_i4bf8.cpp | 6 - .../test_gemm_aquant_basic_i4f32bf8.cpp | 6 - .../test_gemm_aquant_basic_i4f32fp8.cpp | 6 - .../test_gemm_aquant_basic_i4fp8.cpp | 6 - .../test_gemm_aquant_utils.hpp | 243 ----- .../gemm_block_scale/test_gemm_quant_base.hpp | 179 ++++ .../test_gemm_quant_fixtures.hpp | 919 ++++++++++++++++++ .../test_gemm_quant_typed.cpp | 64 ++ .../test_gemm_quant_ut_cases.inc | 28 + .../test_run_gemm_aquant_example.inc | 616 ------------ 39 files changed, 1555 insertions(+), 1056 deletions(-) delete mode 100644 include/ck_tile/ops/gemm_group_quant.hpp create mode 100644 include/ck_tile/ops/gemm_quant.hpp rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/block/block_universal_gemm_as_aquant_bs_cr.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/block/block_universal_gemm_as_bs_bquant_cr.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/kernel/gemm_quant_kernel.hpp (97%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/kernel/grouped_gemm_quant_kernel.hpp (99%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp (99%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp (99%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_group_quant_utils.hpp (100%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/gemm_quant_pipeline_problem.hpp (87%) rename include/ck_tile/ops/{gemm_group_quant => gemm_quant}/pipeline/tile_gemm_quant_traits.hpp (80%) delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_bf8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_fp8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4bf8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32bf8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32fp8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4fp8.cpp delete mode 100644 test/ck_tile/gemm_block_scale/test_gemm_aquant_utils.hpp create mode 100644 test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp create mode 100644 test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp create mode 100644 test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp create mode 100644 test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc delete mode 100644 test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc diff --git a/CHANGELOG.md b/CHANGELOG.md index 6dd06195c9..f21795012d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj * Added benchmarking support for tile engine GEMM Multi D. * Added block scaling support in CK_TILE GEMM, allowing flexible use of quantization matrices from either A or B operands. * Added the row-wise column-wise quantization for CK_TILE GEMM & CK_TILE Grouped GEMM. +* Added tensor-wise quantization for CK_TILE GEMM ### Optimized diff --git a/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp b/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp index 83542e76f1..409bb173a1 100644 --- a/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp +++ b/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp @@ -13,7 +13,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" +#include "ck_tile/ops/gemm_quant.hpp" #include "ck_tile/host.hpp" #include "quant_grouped_gemm.hpp" @@ -65,15 +65,15 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, constexpr auto memory_operation = memory_operation_.value; constexpr bool transpose_c = false; - using QuantGemmProblem = ck_tile::GemmRowColQuantPipelineProblem; + using QuantGemmProblem = ck_tile::GemmRowColTensorQuantPipelineProblem; using GemmPipeline = typename PipelineTypeTraits< GemmConfig::Pipeline>::template GemmPipeline; diff --git a/example/ck_tile/38_block_scale_gemm/README.md b/example/ck_tile/38_block_scale_gemm/README.md index 9acc4f9bfc..9b2610813c 100644 --- a/example/ck_tile/38_block_scale_gemm/README.md +++ b/example/ck_tile/38_block_scale_gemm/README.md @@ -5,6 +5,7 @@ This folder contains examples of quant GEMMs using the ck_tile tile-programming - AQuant kernel with blocks of A matrix sharing scales: custom GEMM pipeline - BQuant kernel with blocks of B matrix sharing scales: custom GEMM pipeline - Row and Column-wise scaled: scaling implemented in Epilogue +- Tensor-wise scaled: scaling implemented in Epilogue ## build ``` @@ -14,7 +15,6 @@ mkdir build && cd build ../script/cmake-ck-dev.sh ../ # Compile the quant kernels make tile_example_gemm_quant_basic -j -make tile_example_gemm_bquant_basic -j ``` This will result in an executable `build/bin/tile_example_gemm_quant_basic` @@ -37,7 +37,7 @@ args: -warmup number of iterations before benchmark the kernel (default:10) -repeat number of iterations to benchmark the kernel (default:100) -timer gpu:gpu timer, cpu:cpu timer (default:gpu) - -quant_mode Which quant method to use (aquant, rowcol) + -quant_mode Which quant method to use (aquant, bquant, tensor, rowcol) ``` User need to select correct mapping of config for each quant mode: diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp index 79c6cca6cb..91f799f194 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp @@ -66,19 +66,21 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str constexpr auto tail_number_v = tail_number_.value; constexpr bool transpose_c = false; + // row-col and tensor quants use the regular pipeline, A/B quants use their own using PipelineProblem = std::conditional_t< - QuantMode == ck_tile::QuantType::RowColQuant, - ck_tile::GemmRowColQuantPipelineProblem, + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant, + ck_tile::GemmRowColTensorQuantPipelineProblem, std::conditional_t>>; using GemmPipeline = std::conditional_t< - QuantMode == ck_tile::QuantType::RowColQuant, + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant, ck_tile::GemmPipelineAgBgCrCompV3, std::conditional_t, @@ -241,10 +244,18 @@ int run_gemm_example(int argc, char* argv[]) ck_tile::QuantType::RowColQuant>( a_layout, b_layout, argc, argv); } + else if(quant_mode == "tensor") + { + return run_gemm_example_prec_type, + TypeConfig, + 128, + ck_tile::QuantType::TensorQuant>( + a_layout, b_layout, argc, argv); + } else { throw std::runtime_error( - "Unsupported quantization mode! Use 'aquant', 'bquant' or 'rowcol'"); + "Unsupported quantization mode! Use 'aquant', 'bquant', 'tensor' or 'rowcol'"); } } else if(data_type == "bf8") @@ -276,10 +287,18 @@ int run_gemm_example(int argc, char* argv[]) ck_tile::QuantType::RowColQuant>( a_layout, b_layout, argc, argv); } + else if(quant_mode == "tensor") + { + return run_gemm_example_prec_type, + TypeConfig, + 128, + ck_tile::QuantType::TensorQuant>( + a_layout, b_layout, argc, argv); + } else { throw std::runtime_error( - "Unsupported quantization mode! Use 'aquant', 'bquant' or 'rowcol'"); + "Unsupported quantization mode! Use 'aquant', 'bquant', 'tensor' or 'rowcol'"); } } else if(data_type == "i4fp8") diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp index ccf07460fa..e5313d8aaf 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp @@ -9,7 +9,7 @@ #include "ck_tile/host/kernel_launch.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" +#include "ck_tile/ops/gemm_quant.hpp" template constexpr ck_tile::index_t get_k_warp_tile() @@ -241,7 +241,7 @@ auto create_args(int argc, char* argv[]) .insert("init", "0", "0:random, 1:linear, 2:constant(1)") .insert("flush_cache", "true", "flush cache before running the kernel, defaults to true") .insert("rotating_count", "1", "rotating count, defaults to 1") - .insert("quant_mode", "aquant", "Choose aquant (default), bquant or rowcol"); + .insert("quant_mode", "aquant", "Choose aquant (default), bquant, tensor or rowcol"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); diff --git a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc index 0f45811ff3..8e9456e973 100644 --- a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc +++ b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc @@ -119,11 +119,7 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, } std::cout << " Acc_Type = " << DataTypeTraits::name << " C_Type = " << DataTypeTraits::name - << " QuantMode = " - << (QuantMode == ck_tile::QuantType::AQuantGrouped - ? "AQuantGrouped" - : (QuantMode == ck_tile::QuantType::BQuantGrouped ? "BQuantGrouped" - : "RowColQuant")) + << " QuantMode = " << quant_type_to_string(QuantMode) << " PreshuffleQuant = " << (GemmConfig::PreshuffleQuant ? "true" : "false") << " : " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; @@ -183,10 +179,11 @@ int run_gemm_example_with_layouts(int argc, AQK = 0; // No A quantization BQK = K / QuantGroupSize; // Group quantization: BQK = K / GroupSize } - else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { - AQK = 1; // Row quantization: tensor shape [M, 1] - BQK = N; // Column quantization: tensor shape [1, N] + AQK = 1; // Row quantization: tensor shape [M, 1] or [1] + BQK = 1; // Column quantization: tensor shape [1, N] or [1] } else { @@ -227,6 +224,11 @@ int run_gemm_example_with_layouts(int argc, stride_AQ = ck_tile::get_default_stride(M, 1, stride_AQ, is_row_major(aq_layout)); stride_BQ = ck_tile::get_default_stride(1, N, stride_BQ, is_row_major(bq_layout)); } + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) + { + stride_AQ = 1; // Tensor quantization: tensor shape [1] + stride_BQ = 1; // Tensor quantization: tensor shape [1] + } ck_tile::HostTensor a_m_k( ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout))); @@ -237,28 +239,30 @@ int run_gemm_example_with_layouts(int argc, // Create AQ tensor with appropriate shape std::unique_ptr> aq_tensor_ptr = nullptr; - if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped) + if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || + QuantMode == ck_tile::QuantType::RowColQuant) { aq_tensor_ptr = std::make_unique>( ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); } - else if(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) { aq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); + ck_tile::host_tensor_descriptor(1, 1, stride_AQ, is_row_major(aq_layout))); } - // Create BQ tensor only for RowColQuant mode + // Create BQ tensor with appropriate shape std::unique_ptr> bq_tensor_ptr = nullptr; - if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped) + if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || + QuantMode == ck_tile::QuantType::RowColQuant) { bq_tensor_ptr = std::make_unique>( ck_tile::host_tensor_descriptor(BQK, N, stride_BQ, is_row_major(bq_layout))); } - else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) { bq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(1, N, stride_BQ, is_row_major(bq_layout))); + ck_tile::host_tensor_descriptor(1, 1, stride_BQ, is_row_major(bq_layout))); } std::random_device rd; @@ -282,7 +286,7 @@ int run_gemm_example_with_layouts(int argc, *bq_tensor_ptr); ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(a_m_k); } - else + else if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped) { if constexpr(std::is_same_v) { @@ -296,12 +300,15 @@ int run_gemm_example_with_layouts(int argc, ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( *aq_tensor_ptr); ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(b_k_n); - - if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) - { - ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( - *bq_tensor_ptr); - } + } + else + { + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(a_m_k); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(b_k_n); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( + *aq_tensor_ptr); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( + *bq_tensor_ptr); } } else if(init_method == 1) @@ -343,7 +350,8 @@ int run_gemm_example_with_layouts(int argc, std::unique_ptr aq_dev_buf_ptr = nullptr; if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { aq_dev_buf_ptr = std::make_unique(aq_tensor_ptr->get_element_space_size_in_bytes()); @@ -351,14 +359,16 @@ int run_gemm_example_with_layouts(int argc, std::unique_ptr bq_dev_buf_ptr = nullptr; if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { bq_dev_buf_ptr = std::make_unique(bq_tensor_ptr->get_element_space_size_in_bytes()); } if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { if constexpr(GemmConfig::PreshuffleQuant) { @@ -398,7 +408,8 @@ int run_gemm_example_with_layouts(int argc, c_m_n_dev_result.SetZero(); if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { bq_dev_buf_ptr->ToDevice(bq_tensor_ptr->data()); } @@ -412,15 +423,9 @@ int run_gemm_example_with_layouts(int argc, CLayout, QuantGroupSize, QuantMode>(a_m_k_dev_buf, - (QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) - ? aq_dev_buf_ptr.get() - : nullptr, + aq_dev_buf_ptr.get(), b_k_n_dev_buf, - (QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) - ? bq_dev_buf_ptr.get() - : nullptr, + bq_dev_buf_ptr.get(), c_m_n_dev_buf, M, N, @@ -467,7 +472,7 @@ int run_gemm_example_with_layouts(int argc, QuantGroupSize, false>(a_m_k, *bq_tensor_ptr, b_k_n, c_m_n_host_ref); } - else + else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) { ck_tile::reference_gemm_rowcol_quant( a_m_k, *aq_tensor_ptr, b_k_n, *bq_tensor_ptr, c_m_n_host_ref); } + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) + { + ck_tile::reference_gemm_tensor_quant( + a_m_k, *aq_tensor_ptr, b_k_n, *bq_tensor_ptr, c_m_n_host_ref); + } const float max_accumulated_value = *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index c7c4702e22..a3620453b4 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -158,4 +158,7 @@ CK_TILE_DEVICE auto load_tile_raw(T& /*null_tile*/, const null_tile_window +concept IsLoadableTile = requires { load_tile(std::declval()); }; + } // namespace ck_tile diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index d9379b4420..90f68f7e2e 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -180,10 +180,6 @@ CK_TILE_HOST void reference_gemm_rowcol_quant(const HostTensor& a_m_k else v_b = fp32_val.lo; } - else if constexpr(std::is_same_v) - { - v_b = fp8_to_float_raw(b_element_op(b_k_n(k, n))); - } else { v_b = ck_tile::type_convert(b_element_op(b_k_n(k, n))); @@ -198,7 +194,57 @@ CK_TILE_HOST void reference_gemm_rowcol_quant(const HostTensor& a_m_k }; make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); - std::cout << std::endl; +} + +template +CK_TILE_HOST void reference_gemm_tensor_quant(const HostTensor& a_m_k, + const HostTensor& aq_1_1, + const HostTensor& b_k_n, + const HostTensor& bq_1_1, + HostTensor& c_m_n, + const AElementOp& a_element_op = {}, + const BElementOp& b_element_op = {}, + const ACCElementOp& acc_element_op = {}) +{ + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v && std::is_same_v); + const std::size_t M = a_m_k.get_length(0); + const std::size_t N = b_k_n.get_length(1); + const std::size_t K = a_m_k.get_length(1); + + auto f_mn = [&](auto m, auto n) { + // Init accumulator + AccDataType v_acc = 0; + // Get scale for A and scale for B + const AccDataType a_scale = ck_tile::type_convert(aq_1_1(0, 0)); + const AccDataType b_scale = ck_tile::type_convert(bq_1_1(0, 0)); + + // Compute the dot product + for(std::size_t k = 0; k < K; ++k) + { + AccDataType v_a = ck_tile::type_convert(a_element_op(a_m_k(m, k))); + AccDataType v_b = ck_tile::type_convert(b_element_op(b_k_n(k, n))); + + v_acc += v_a * v_b; + } + + v_acc = v_acc * a_scale * b_scale; + + c_m_n(m, n) = ck_tile::type_convert(acc_element_op(v_acc)); + }; + + make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); } template && std::is_same_v) { - constexpr auto step = SFC::get_forward_step(iAccess); + // No scaling needed - this is a no-op + } + // Check if scales are scalar AccDataType + else if constexpr(std::is_same_v && + std::is_same_v) + { + // Handle scalar scales + const AccDataType scale_m = scale_m_window; + const AccDataType scale_n = scale_n_window; + tile_elementwise_inout([&](auto& element) { element = element * scale_m * scale_n; }, + lds_tile); + } + // Otherwise, assume they are tile windows that can be loaded + else + { + // Load tiles + const auto scale_m_tile = load_tile(scale_m_window); + const auto scale_n_tile = load_tile(scale_n_window); - move_tile_window(scale_m_window, {step.at(number<0>{}), step.at(number<1>{})}); - move_tile_window(scale_n_window, {step.at(number<0>{}), step.at(number<1>{})}); + // Compute element-wise product in-place i.e. lds_tile = lds_tile * scale_m * scale_n + tile_elementwise_inout( + element_wise::MultiDMultiply{}, lds_tile, lds_tile, scale_m_tile, scale_n_tile); + + // Move scale windows + constexpr index_t num_access = SFC::get_num_of_access(); + if constexpr(iAccess != num_access - 1) + { + constexpr auto step = SFC::get_forward_step(iAccess); + + move_tile_window(scale_m_window, {step.at(number<0>{}), step.at(number<1>{})}); + move_tile_window(scale_n_window, {step.at(number<0>{}), step.at(number<1>{})}); + } } } @@ -452,6 +471,8 @@ struct CShuffleEpilogue // Optional scales (must share the same distribution to match per-thread indexing) constexpr bool has_scales = !std::is_same::value && !std::is_same::value; + constexpr bool has_scalar_scales = + std::is_same_v && std::is_same_v; // Tiles to hold row/col scales when present using SMType = typename GetDataType>::type; @@ -462,8 +483,11 @@ struct CShuffleEpilogue // Build windows only if scales are provided auto scale_m_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scales && !has_scalar_scales) { + static_assert( + IsLoadableTile, + "ScaleM must be a loadable tile"); return make_tile_window(scale_m, dram_tile_distribution); } else @@ -472,8 +496,11 @@ struct CShuffleEpilogue } }(); auto scale_n_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scales && !has_scalar_scales) { + static_assert( + IsLoadableTile, + "ScaleN must be a loadable tile"); return make_tile_window(scale_n, dram_tile_distribution); } else @@ -489,7 +516,7 @@ struct CShuffleEpilogue merge_sequences(sequence<1, NRepeat>{}, c_warp_y_lengths)); // If scales provided, load them with identical distribution - if constexpr(has_scales) + if constexpr(has_scales && IsLoadableTile && IsLoadableTile) { sm_tile = load_tile(scale_m_window); // row scales in permuted layout sn_tile = load_tile(scale_n_window); // col scales in permuted layout @@ -504,7 +531,11 @@ struct CShuffleEpilogue auto emit = [&](index_t out_idx, index_t src_row) { AccDataType v = shuffle_acc.get_thread_buffer()[base + src_row]; - if constexpr(has_scales) + if constexpr(has_scalar_scales) + { + v = static_cast(v * scale_m * scale_n); + } + else if constexpr(has_scales) { // same linear index mapping on the permuted distribution const auto s_m = static_cast(sm_tile.get_thread_buffer()[out_idx]); @@ -595,10 +626,19 @@ struct CShuffleEpilogue number{}); constexpr bool has_scales = - !std::is_same::value && !std::is_same::value; + !std::is_same_v && !std::is_same_v; + constexpr bool has_scalar_scales = + std::is_same_v && std::is_same_v; auto scale_m_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scalar_scales) { + return scale_m; + } + else if constexpr(has_scales) + { + static_assert( + IsLoadableTile, + "ScaleM must be a loadable tile"); return make_tile_window(scale_m, lds_tile.get_tile_distribution()); } else @@ -607,8 +647,15 @@ struct CShuffleEpilogue } }(); auto scale_n_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scalar_scales) { + return scale_n; + } + else if constexpr(has_scales) + { + static_assert( + IsLoadableTile, + "ScaleN must be a loadable tile"); return make_tile_window(scale_n, lds_tile.get_tile_distribution()); } else diff --git a/include/ck_tile/ops/gemm_group_quant.hpp b/include/ck_tile/ops/gemm_group_quant.hpp deleted file mode 100644 index 94b5ab8c3b..0000000000 --- a/include/ck_tile/ops/gemm_group_quant.hpp +++ /dev/null @@ -1,21 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp" -#include "ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.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/gemm_quant.hpp b/include/ck_tile/ops/gemm_quant.hpp new file mode 100644 index 0000000000..9f90050899 --- /dev/null +++ b/include/ck_tile/ops/gemm_quant.hpp @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp" +#include "ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp" +#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.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/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp rename to include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp rename to include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp similarity index 97% rename from include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp rename to include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp index 13fa0b8dfb..82bf75a9e3 100644 --- a/include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp @@ -12,7 +12,7 @@ #include "ck_tile/core/numeric/integer.hpp" #include "ck_tile/core/numeric/math.hpp" #include "ck_tile/host/concat.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp" namespace ck_tile { @@ -330,7 +330,6 @@ struct QuantGemmKernel } } - // NOTE: no kernel currently uses BQuant like this: if constexpr(kQuantType == QuantType::BQuantGrouped) { static_assert(std::is_same_v); @@ -890,6 +889,7 @@ struct QuantGemmKernel * @param a_ptr input A pointer * @param b_ptr input B pointer * @param aq_ptr input AQ pointer + * @param bq_ptr input BQ pointer * @param c_ptr output C pointer * @param smem_ptr_0 The start memory pointer of the shared memory block. * @param kargs GEMM kernel arguments @@ -938,7 +938,8 @@ struct QuantGemmKernel return GemmPipeline{}.template operator()( a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr_0); } - else if constexpr(kQuantType == QuantType::RowColQuant) + else if constexpr(kQuantType == QuantType::RowColQuant || + kQuantType == QuantType::TensorQuant) { return GemmPipeline{}.template operator()( a_block_window, b_block_window, num_loop, smem_ptr_0); @@ -964,6 +965,18 @@ struct QuantGemmKernel aq_block_window, bq_block_window); } + else if constexpr(kQuantType == QuantType::TensorQuant) + { + // TODO: why doesn't readfirstlane work here? + // const AccDataType aq_scale = + // __builtin_amdgcn_readfirstlane(type_convert(*aq_ptr)); + // const AccDataType bq_scale = + // __builtin_amdgcn_readfirstlane(type_convert(*bq_ptr)); + const AccDataType aq_scale = type_convert(*aq_ptr); + const AccDataType bq_scale = type_convert(*bq_ptr); + EpiloguePipeline{}( + c_block_window, c_block_tile, c_block_window, smem_ptr_0, aq_scale, bq_scale); + } } CK_TILE_DEVICE void operator()(QuantGemmKernelArgs kargs) const diff --git a/include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp rename to include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp index 925ea42678..07c45117e2 100644 --- a/include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp @@ -9,7 +9,7 @@ #include "ck_tile/host/stream_utils.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp" #include "ck_tile/host.hpp" #include diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp index 5ce4268dca..24254013a4 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp @@ -9,7 +9,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/numeric/math.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp index 8f191f0f94..c27fbf5b50 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp @@ -9,7 +9,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp similarity index 87% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp index a2cef2d994..d49204c64d 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp @@ -168,17 +168,18 @@ template -using GemmRowColQuantPipelineProblem = GemmQuantPipelineProblemBase; +using GemmRowColTensorQuantPipelineProblem = + GemmQuantPipelineProblemBase; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp similarity index 80% rename from include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp index f505efe4e0..e97eeffb9b 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp @@ -12,9 +12,22 @@ enum struct QuantType : std::uint16_t { AQuantGrouped = 0, BQuantGrouped = 1, - RowColQuant = 2 + RowColQuant = 2, + TensorQuant = 3 }; +std::string quant_type_to_string(QuantType quant_type) +{ + switch(quant_type) + { + case QuantType::AQuantGrouped: return "AQuantGrouped"; + case QuantType::BQuantGrouped: return "BQuantGrouped"; + case QuantType::RowColQuant: return "RowColQuant"; + case QuantType::TensorQuant: return "TensorQuant"; + default: return "Unknown"; + } +} + template ; - bool result = run_cshuffle_epilogue_test(); - EXPECT_TRUE(result) << "Basic CShuffleEpilogue test failed"; + auto result = run_cshuffle_epilogue_test(ScaleType::None); + EXPECT_FLOAT_EQ(result[0], 2.0F) << "Basic CShuffleEpilogue test failed"; } TEST_F(CShuffleEpilogueTest, BasicHalfTestWithScale) @@ -73,8 +73,45 @@ TEST_F(CShuffleEpilogueTest, BasicHalfTestWithScale) NPerXdl, KPerXdl>; - bool result = run_cshuffle_epilogue_test(true); - EXPECT_TRUE(result) << "Scale CShuffleEpilogue test failed"; + auto result = + run_cshuffle_epilogue_test(ScaleType::RowCol); + EXPECT_FLOAT_EQ(result[0], 2.0F) << "RowCol CShuffleEpilogue test failed: first element not 2"; + EXPECT_FLOAT_EQ(result[1], 4.0F) + << "RowCol CShuffleEpilogue test failed: second element not 2*2"; +} + +TEST_F(CShuffleEpilogueTest, BasicHalfTestWithTensorScale) +{ + // Basic test configuration with half_t data types + using ADataType = ck_tile::half_t; + using BDataType = ck_tile::half_t; + using AccDataType = float; + using ODataType = ck_tile::half_t; + + constexpr index_t kMPerBlock = 256; + constexpr index_t kNPerBlock = 256; + constexpr index_t MWave = 2; + constexpr index_t NWave = 2; + constexpr index_t MPerXdl = 32; + constexpr index_t NPerXdl = 32; + constexpr index_t KPerXdl = 8; + + using TestProblem = SimpleCShuffleEpilogueProblem; + + auto result = + run_cshuffle_epilogue_test(ScaleType::Tensor); + EXPECT_FLOAT_EQ(result[0], 4.0F) + << "TensorScale CShuffleEpilogue test failed: first element not 2*2=4"; } int main(int argc, char** argv) diff --git a/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp b/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp index c23957d802..01e6c91c7c 100644 --- a/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp +++ b/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp @@ -19,8 +19,15 @@ namespace ck_tile { +enum class ScaleType +{ + None, + RowCol, + Tensor +}; + // Simple test kernel to invoke the CShuffleEpilogue -template +template __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __restrict__ output_data, float* m_scale, float* n_scale) @@ -61,7 +68,7 @@ __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __res auto empty_ds = make_tuple(); // Call the epilogue - if constexpr(UseScale) + if constexpr(Scale == ScaleType::RowCol) { const auto m_scale_window = make_tile_window( make_naive_tensor_view( @@ -75,6 +82,10 @@ __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __res {0, 0}); Epilogue{}(output_tile_window, acc_tile, empty_ds, smem, m_scale_window, n_scale_window); } + else if constexpr(Scale == ScaleType::Tensor) + { + Epilogue{}(output_tile_window, acc_tile, empty_ds, smem, *m_scale, *n_scale); + } else { Epilogue{}(output_tile_window, acc_tile, empty_ds, smem); @@ -113,7 +124,7 @@ using SimpleCShuffleEpilogueProblem = memory_operation_enum::set>; template -bool run_cshuffle_epilogue_test(bool use_scale = false) +auto run_cshuffle_epilogue_test(ScaleType scale = ScaleType::None) { using ODataType = typename Problem::ODataType; @@ -142,7 +153,7 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) dim3 gridSize(1, 1, 1); dim3 blockSize(kBlockSize, 1, 1); - if(use_scale) + if(scale == ScaleType::RowCol) { float* m_scale; float* n_scale; @@ -155,12 +166,25 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) hipMemcpy(m_scale, h_m_scale.data(), M * sizeof(float), hipMemcpyHostToDevice)); HIP_CHECK_ERROR( hipMemcpy(n_scale, h_n_scale.data(), N * sizeof(float), hipMemcpyHostToDevice)); - test_cshuffle_epilogue_kernel + test_cshuffle_epilogue_kernel + <<>>(device_output, m_scale, n_scale); + } + else if(scale == ScaleType::Tensor) + { + float* m_scale; + float* n_scale; + std::vector h_m_scale(1, 2.0F); + std::vector h_n_scale(1, 1.0F); + HIP_CHECK_ERROR(hipMalloc(&m_scale, sizeof(float))); + HIP_CHECK_ERROR(hipMalloc(&n_scale, sizeof(float))); + HIP_CHECK_ERROR(hipMemcpy(m_scale, h_m_scale.data(), sizeof(float), hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemcpy(n_scale, h_n_scale.data(), sizeof(float), hipMemcpyHostToDevice)); + test_cshuffle_epilogue_kernel <<>>(device_output, m_scale, n_scale); } else { - test_cshuffle_epilogue_kernel + test_cshuffle_epilogue_kernel <<>>(device_output, nullptr, nullptr); } @@ -172,20 +196,10 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) HIP_CHECK_ERROR(hipMemcpy( host_output.data(), device_output, output_size * sizeof(ODataType), hipMemcpyDeviceToHost)); - // Basic verification - just check that output has a 2, and 4 if using scaling - bool has_2 = - type_convert(host_output[0]) > 1.9F && type_convert(host_output[0]) < 2.1F; - bool scale_has_4 = true; - if(use_scale) - { - scale_has_4 = type_convert(host_output[1]) > 3.9F && - type_convert(host_output[1]) < 4.1F; - } - // Cleanup HIP_CHECK_ERROR(hipFree(device_output)); - return has_2 && scale_has_4; + return host_output; } } // namespace ck_tile diff --git a/test/ck_tile/gemm_block_scale/CMakeLists.txt b/test/ck_tile/gemm_block_scale/CMakeLists.txt index 847ab88644..93a13ba5af 100644 --- a/test/ck_tile/gemm_block_scale/CMakeLists.txt +++ b/test/ck_tile/gemm_block_scale/CMakeLists.txt @@ -6,14 +6,9 @@ endif() list(APPEND TEST_GEMM_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion=0) if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95") - set(TEST_GEMM_NAME test_tile_gemm_aquant_basic) - set(QUANT_TYPES fp8 bf8 i4fp8 i4bf8 i4f32fp8 i4f32bf8) - - foreach(QUANT_TYPE ${QUANT_TYPES}) - add_gtest_executable(${TEST_GEMM_NAME}_${QUANT_TYPE} test_gemm_aquant_basic_${QUANT_TYPE}.cpp) - target_compile_options(${TEST_GEMM_NAME}_${QUANT_TYPE} PRIVATE ${TEST_GEMM_COMPILE_OPTIONS}) - endforeach() - + # Typed Test Suite for GEMM Quantization + add_gtest_executable(test_tile_gemm_quant_typed test_gemm_quant_typed.cpp) + target_compile_options(test_tile_gemm_quant_typed PRIVATE ${TEST_GEMM_COMPILE_OPTIONS}) else() message(DEBUG "Skipping ck_tile quant gemm tests for current target") endif() diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_bf8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_bf8.cpp deleted file mode 100644 index 9c4277d879..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_bf8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("bf8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_fp8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_fp8.cpp deleted file mode 100644 index b0cf55be6f..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_fp8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("fp8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4bf8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4bf8.cpp deleted file mode 100644 index fd80bf2b06..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4bf8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("i4bf8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32bf8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32bf8.cpp deleted file mode 100644 index fe8c9c5000..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32bf8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("i4f32bf8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32fp8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32fp8.cpp deleted file mode 100644 index a319d9c2ad..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4f32fp8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("i4f32fp8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4fp8.cpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4fp8.cpp deleted file mode 100644 index ceb8760435..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_basic_i4fp8.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#include "test_run_gemm_aquant_example.inc" - -int main() { return run_gemm_combinations("i4fp8"); } diff --git a/test/ck_tile/gemm_block_scale/test_gemm_aquant_utils.hpp b/test/ck_tile/gemm_block_scale/test_gemm_aquant_utils.hpp deleted file mode 100644 index 83a9e57878..0000000000 --- a/test/ck_tile/gemm_block_scale/test_gemm_aquant_utils.hpp +++ /dev/null @@ -1,243 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/epilogue.hpp" -#include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" - -#define CK_TILE_PIPELINE_PREFILL 1 -#define CK_TILE_PIPELINE_DECODE 2 -#define CK_TILE_PIPELINE_PRESHUFFLEQUANT 3 - -template -constexpr ck_tile::index_t get_k_warp_tile() -{ -#if defined(CK_GFX950_SUPPORT) - constexpr bool is_8bit_float = - std::is_same_v || std::is_same_v; - if constexpr(M_Warp_Tile == 32) - return is_8bit_float ? 64 : 16; - else - return is_8bit_float ? 128 : 32; -#else - if constexpr(M_Warp_Tile == 32) - return 16; - else - return 32; -#endif -} - -template -auto calculate_rtol_atol(const ck_tile::index_t K, - const ck_tile::index_t kbatch, - const float max_accumulated_value) -{ - using ComputeType = - std::conditional_t; - // Calculate thresholds - const auto rtol = ck_tile::get_relative_threshold( - ck_tile::integer_divide_ceil(K, kbatch)); - const auto atol = ck_tile::get_absolute_threshold( - max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); - // Calculate error due to split_k accumulation - const auto rtol_split_k = - ck_tile::get_relative_threshold(kbatch); - const auto atol_split_k = ck_tile::get_absolute_threshold( - max_accumulated_value, kbatch); - // Use higher threshold - return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); -} - -class ArgumentsNotSupportedException : public std::logic_error -{ - public: - explicit ArgumentsNotSupportedException(const std::string& message) : logic_error(message) {} -}; - -struct GemmConfigBase -{ - static constexpr bool kPadM = false; - static constexpr bool kPadN = false; - static constexpr bool kPadK = false; - - static constexpr bool PermuteA = false; - static constexpr bool PermuteB = false; - - static constexpr bool TransposeC = false; - static constexpr bool UseStructuredSparsity = false; - - static constexpr int kBlockPerCu = 1; - static constexpr ck_tile::index_t TileParitionerGroupNum = 8; - static constexpr ck_tile::index_t TileParitionerM01 = 4; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; - static constexpr ck_tile::index_t NumWaveGroups = 1; - static constexpr bool PreshuffleQuant = false; - static constexpr bool DoubleSmemBuffer = true; -}; - -template -struct GemmConfigDecode : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 16; - static constexpr ck_tile::index_t N_Tile = 64; - static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile(); - - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_DECODE; -}; - -template -struct GemmConfigPrefill : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 128; - static constexpr ck_tile::index_t N_Tile = 128; - static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile(); - - static constexpr int kBlockPerCu = 2; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_PREFILL; -}; - -template -struct GemmConfigPreshuffleQuant : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 16; - static constexpr ck_tile::index_t N_Tile = 64; - static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = - get_k_from_preshuffled_warp_tile(); - - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_PRESHUFFLEQUANT; - static constexpr bool PreshuffleQuant = true; -}; - -template -struct GemmQuantTypeConfig -{ - using ADataType = ADataType_; - using QDataType = QDataType_; - using BDataType = BDataType_; - using AccDataType = float; - using CDataType = CDataType_; -}; - -template -struct DataTypeTraits; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp32"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp64"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "int32"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp16"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "bf16"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp8"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "bf8"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "pk_int4_t"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "int8"; -}; - -auto create_args(int argc, char* argv[]) -{ - ck_tile::ArgParser arg_parser; - arg_parser.insert("m", "3840", "m dimension") - .insert("n", "4096", "n dimension") - .insert("k", "2048", "k dimension") - .insert("a_layout", "R", "A tensor data layout - Row by default") - .insert("aq_layout", "R", "Aq tensor data layout - Row by default") - .insert("b_layout", "C", "B tensor data layout - Column by default") - .insert("c_layout", "R", "C tensor data layout - Row by default") - .insert("stride_a", "0", "Tensor A stride") - .insert("stride_q", "0", "Tensor AQ stride") - .insert("stride_b", "0", "Tensor B stride") - .insert("stride_c", "0", "Tensor C stride") - .insert("v", "2", "0. No validation, 1. Validation on CPU, 2. Validation on GPU") - .insert("prec", "i4fp8", "data type. fp8/bf8/i4fp8/i4bf8/i4f32fp8/i4f32bf8") - .insert("warmup", "50", "number of iterations before benchmark the kernel") - .insert("repeat", "100", "number of iterations to benchmark the kernel") - .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer") - .insert("split_k", "1", "splitK value") - .insert("init", "0", "0:random, 1:linear, 2:constant(1)") - .insert("persistent", "0", "0:non-persistent, 1:persistent") - .insert("as_br_cr", "false", "Choose between as_br_cr and as_bs_cr"); - - bool result = arg_parser.parse(argc, argv); - return std::make_tuple(result, arg_parser); -} - -// host API -float gemm_calc_aquant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s); diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp new file mode 100644 index 0000000000..ed3231d140 --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/check_err.hpp" +#include "ck_tile/host/reference/reference_gemm.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm_quant.hpp" + +// Forward declarations for quant type-specific implementations +template +struct QuantTypeTraits; + +// Base class for common quant gemm functionality +template +class TestCkTileGemmQuantBase : public ::testing::Test +{ + protected: + using ALayout = std::tuple_element_t<0, Tuple>; + using BLayout = std::tuple_element_t<1, Tuple>; + using CLayout = std::tuple_element_t<2, Tuple>; + using ADataType = std::tuple_element_t<3, Tuple>; + using BDataType = std::tuple_element_t<4, Tuple>; + using QDataType = std::tuple_element_t<5, Tuple>; + using CDataType = std::tuple_element_t<6, Tuple>; + static constexpr auto QuantType = std::tuple_element_t<7, Tuple>::value; + using GemmConfig = std::tuple_element_t<8, Tuple>; + static constexpr uint32_t QuantGroupSize = std::tuple_element_t<9, Tuple>::value; + using AccDataType = float; // accumulate always in float + + // Get the quant-type specific data types from traits + using QuantTraits = QuantTypeTraits; + using ComputeDataType = typename QuantTraits::template ComputeDataType; + + static constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; + static constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; + static constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; + + static constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; + static constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; + static constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; + + static constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; + static constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; + static constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; + + public: + void SetUp() override { static_cast(this)->SetUpQuantTypeSpecific(); } + + void TearDown() override { static_cast(this)->TearDownQuantTypeSpecific(); } + + // Common test execution logic + void invoke_quant_gemm(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s) + { + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + constexpr bool kPreshuffle = false; + + using CodegenGemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using CodegenGemmTraits = ck_tile::TileGemmQuantTraits; + + // Let the derived class create the appropriate pipeline and epilogue + static_cast(this) + ->template run_quant_gemm_impl( + args, s); + } + + void RunTest(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + // Generate test data and run the kernel + static_cast(this)->run_test_with_validation(M, N, K); + } + + // Helper function to check layout + template + static constexpr auto is_row_major(Layout) + { + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; + } + + // Tolerance calculation function for validation + template + auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) + { + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = + ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); + } +}; + +// Define generic QuantTypeTraits template (will be specialized) +template +struct QuantTypeTraits +{ + static_assert(QT == ck_tile::QuantType::AQuantGrouped || + QT == ck_tile::QuantType::BQuantGrouped || + QT == ck_tile::QuantType::RowColQuant || + QT == ck_tile::QuantType::TensorQuant, + "Unsupported quantization type"); +}; + +// Specialization for AQuantGrouped +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = BDataType; // For AQuant, compute type is BDataType + + static constexpr const char* name = "aquant"; +}; + +// Specialization for BQuantGrouped +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For BQuant, compute type is ADataType + + static constexpr const char* name = "bquant"; +}; + +// Specialization for RowColQuant +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For RowColQuant, compute type is ADataType + + static constexpr const char* name = "rowcol"; +}; + +// Specialization for TensorQuant +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For TensorQuant, compute type is ADataType + + static constexpr const char* name = "tensor"; +}; diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp new file mode 100644 index 0000000000..5fc6b2f15c --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp @@ -0,0 +1,919 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "test_gemm_quant_base.hpp" +#include "ck_tile/host/permute_pk_int4.hpp" + +struct GemmConfigBase +{ + static constexpr bool kPadM = false; + static constexpr bool kPadN = false; + static constexpr bool kPadK = false; + + static constexpr bool PermuteA = false; + static constexpr bool PermuteB = false; + + static constexpr bool TransposeC = false; + static constexpr bool UseStructuredSparsity = false; + + static constexpr int kBlockPerCu = 1; + static constexpr ck_tile::index_t TileParitionerGroupNum = 8; + static constexpr ck_tile::index_t TileParitionerM01 = 4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; + static constexpr ck_tile::index_t NumWaveGroups = 1; + static constexpr bool PreshuffleQuant = false; + static constexpr bool DoubleSmemBuffer = false; + + // Default GEMM tile sizes for tests + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 256; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; +}; + +template +class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + // AQuant-specific data generation + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // AQuant uses grouped quantization for A matrix + const ck_tile::index_t AQK = ck_tile::integer_divide_ceil(K, QuantGroupSize); + const ck_tile::index_t stride_AQ = + ck_tile::get_default_stride(M, AQK, 0, this->is_row_major(ALayout{})); + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor aq_m_aqk( + ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + + // Initialize data with random values + if constexpr(std::is_same_v) + { + ck_tile::FillUniformDistribution{-5.0f, 5.0f}(a_m_k); + } + else + { + ck_tile::FillUniformDistribution{-2.0f, 3.0f}(a_m_k); + } + ck_tile::FillUniformDistribution{-5.0f, 5.0f}(b_k_n); + ck_tile::FillUniformDistribution{-2.0f, 2.0f}(aq_m_aqk); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem aq_m_aqk_dev_buf(aq_m_aqk.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor temp = a_m_k; + ck_tile::permute_vectors_i4x4_b(temp); + a_m_k_dev_buf.ToDevice(temp.data()); + } + else + { + a_m_k_dev_buf.ToDevice(a_m_k.data()); + } + aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + aq_m_aqk_dev_buf.GetDeviceBuffer(), // aq_ptr (scales) + nullptr, // bq_ptr (not used for AQuant) + 1, // k_batch + M, + N, + K, // M, N, K + AQK, // QK_A + 0, // QK_B (not used for AQuant) + stride_A, + stride_B, + stride_C, + stride_AQ, + 0 // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference AQuant implementation + ck_tile::reference_gemm_quant(a_m_k, aq_m_aqk, b_k_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "AQuantGrouped validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "AQuantGrouped - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // AQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = + ck_tile::GemmAQuantPipelineProblem; + + using GemmPipeline = ck_tile::AQuantGemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for AQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// BQuant-specific test fixture +template +class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // BQuant uses grouped quantization for B matrix + const ck_tile::index_t BQK = ck_tile::integer_divide_ceil(K, QuantGroupSize); + const ck_tile::index_t stride_BQ = BQK; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor bq_bqk_n( + ck_tile::host_tensor_descriptor(BQK, N, stride_BQ, this->is_row_major(BLayout{}))); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{0.f, 1.f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(bq_bqk_n); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem bq_bqk_n_dev_buf(bq_bqk_n.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor temp = b_k_n; + ck_tile::permute_vectors_i4x4_b(temp); + b_k_n_dev_buf.ToDevice(temp.data()); + } + else + { + b_k_n_dev_buf.ToDevice(b_k_n.data()); + } + bq_bqk_n_dev_buf.ToDevice(bq_bqk_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + nullptr, // aq_ptr (not used for BQuant) + bq_bqk_n_dev_buf.GetDeviceBuffer(), // bq_ptr (scales) + 1, // k_batch + M, + N, + K, // M, N, K + 0, // QK_A (not used for BQuant) + BQK, // QK_B + stride_A, + stride_B, + stride_C, + 0, + stride_BQ // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference BQuant implementation + ck_tile::reference_gemm_quant(a_m_k, bq_bqk_n, b_k_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "BQuantGrouped validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "BQuantGrouped - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // BQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + + using PipelineProblem = + ck_tile::GemmBQuantPipelineProblem; + + using GemmPipeline = ck_tile::BQuantGemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + false, // transpose_c + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for BQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// RowColQuant-specific test fixture +template +class TestCkTileGemmRowColQuant + : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // RowColQuant uses per-row and per-column scales + const ck_tile::index_t stride_row_scales = 1; + const ck_tile::index_t stride_col_scales = 1; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor row_scales_m(ck_tile::host_tensor_descriptor( + M, 1, stride_row_scales, ck_tile::bool_constant{})); + ck_tile::HostTensor col_scales_n(ck_tile::host_tensor_descriptor( + N, 1, stride_col_scales, ck_tile::bool_constant{})); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(row_scales_m); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(col_scales_n); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem row_scales_dev_buf(row_scales_m.get_element_space_size() * + sizeof(QDataType)); + ck_tile::DeviceMem col_scales_dev_buf(col_scales_n.get_element_space_size() * + sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + row_scales_dev_buf.ToDevice(row_scales_m.data()); + col_scales_dev_buf.ToDevice(col_scales_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + row_scales_dev_buf.GetDeviceBuffer(), // aq_ptr (row scales) + col_scales_dev_buf.GetDeviceBuffer(), // bq_ptr (col scales) + 1, // k_batch + M, + N, + K, // M, N, K + 1, // QK_A (row scales) + 1, // QK_B (col scales) + stride_A, + stride_B, + stride_C, + stride_row_scales, + stride_col_scales // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference RowColQuant implementation + ck_tile::reference_gemm_rowcol_quant( + a_m_k, row_scales_m, b_k_n, col_scales_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "RowColQuant validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "RowColQuant - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // RowColQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = ck_tile::GemmRowColTensorQuantPipelineProblem< + ADataType, + BDataType, + AccDataType, + AccDataType, + CodegenGemmShape, + CodegenGemmTraits, + transpose_c, + ComputeDataType, + ck_tile::GemmPipelineScheduler::Intrawave, + has_hot_loop_v, + tail_number_v>; + + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for RowColQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// TensorQuant-specific test fixture +template +class TestCkTileGemmTensorQuant + : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // TensorQuant uses single scalar scale for each tensor + const ck_tile::index_t stride_scale_a = 1; + const ck_tile::index_t stride_scale_b = 1; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor scale_a( + ck_tile::host_tensor_descriptor(1, 1, stride_scale_a, ck_tile::bool_constant{})); + ck_tile::HostTensor scale_b( + ck_tile::host_tensor_descriptor(1, 1, stride_scale_b, ck_tile::bool_constant{})); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(scale_a); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(scale_b); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem scale_a_dev_buf(scale_a.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem scale_b_dev_buf(scale_b.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + scale_a_dev_buf.ToDevice(scale_a.data()); + scale_b_dev_buf.ToDevice(scale_b.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + scale_a_dev_buf.GetDeviceBuffer(), // aq_ptr (scale A) + scale_b_dev_buf.GetDeviceBuffer(), // bq_ptr (scale B) + 1, // k_batch + M, + N, + K, // M, N, K + 1, // QK_A (tensor scale) + 1, // QK_B (tensor scale) + stride_A, + stride_B, + stride_C, + stride_scale_a, + stride_scale_b // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference TensorQuant implementation + ck_tile::reference_gemm_tensor_quant( + a_m_k, scale_a, b_k_n, scale_b, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "TensorQuant validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "TensorQuant - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // TensorQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = ck_tile::GemmRowColTensorQuantPipelineProblem< + ADataType, + BDataType, + AccDataType, + AccDataType, + CodegenGemmShape, + CodegenGemmTraits, + transpose_c, + ComputeDataType, + ck_tile::GemmPipelineScheduler::Intrawave, + has_hot_loop_v, + tail_number_v>; + + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for TensorQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp new file mode 100644 index 0000000000..1926b7cd0f --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +#include +#include + +#include "test_gemm_quant_fixtures.hpp" + +// Type aliases for readability +using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; +using ColumnMajor = ck_tile::tensor_layout::gemm::ColumnMajor; +using FP8 = ck_tile::fp8_t; +using BF8 = ck_tile::bf8_t; +using Half = ck_tile::half_t; +using PkInt4 = ck_tile::pk_int4_t; +using AQuantGrouped = std::integral_constant; +using BQuantGrouped = std::integral_constant; +using RowColQuant = std::integral_constant; +using TensorQuant = std::integral_constant; +using GroupSize = std::integral_constant; + +// Type combinations for each quantization type +// clang-format off +using AQuantTypes = ::testing::Types< + std::tuple, + std::tuple, + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using BQuantTypes = ::testing::Types< + std::tuple, + std::tuple, + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using RowColQuantTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using TensorQuantTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +// Test suites for each quantization type +TYPED_TEST_SUITE(TestCkTileGemmAQuant, AQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmBQuant, BQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmRowColQuant, RowColQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmTensorQuant, TensorQuantTypes); + +#include "test_gemm_quant_ut_cases.inc" diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc b/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc new file mode 100644 index 0000000000..9b07afa2b3 --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +// AQuant tests +TYPED_TEST(TestCkTileGemmAQuant, AQuantGroupedTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// BQuant tests +TYPED_TEST(TestCkTileGemmBQuant, BQuantGroupedTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// RowColQuant tests +TYPED_TEST(TestCkTileGemmRowColQuant, RowColQuantTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// TensorQuant tests +TYPED_TEST(TestCkTileGemmTensorQuant, TensorQuantTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} diff --git a/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc b/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc deleted file mode 100644 index dbe652ac62..0000000000 --- a/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc +++ /dev/null @@ -1,616 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#include "ck_tile/core/config.hpp" -#include "ck_tile/host.hpp" -#include "test_gemm_aquant_utils.hpp" -#include "ck_tile/host/permute_pk_int4.hpp" - -template -float gemm_calc_aquant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s) -{ - constexpr bool kPadM = false; - constexpr bool kPadN = false; - constexpr bool kPadK = false; - - constexpr int kBlockPerCu = 1; - - static_assert(std::is_same_v); - - constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; - constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; - constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; - - constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; - constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; - constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; - - constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; - constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; - constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; - - using CodegenGemmShape = - ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence>; - - using TilePartitioner = ck_tile::GemmTile1DPartitioner; - - using CodegenGemmTraits = ck_tile::TileGemmQuantTraits; - - using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; - - using BaseGemmPipeline = ck_tile::BaseAQuantGemmPipelineAgBgCrCompV3; - - const ck_tile::index_t K_split = (args.K + K_Tile - 1) / K_Tile * K_Tile; - const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); - const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); - const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); - constexpr bool transposed_warp_gemm = false; - - const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { - constexpr bool has_hot_loop_v = has_hot_loop_.value; - constexpr auto tail_number_v = tail_number_.value; - - using CodegenPipelineProblem = - ck_tile::GemmAQuantPipelineProblem; - using CodegenGemmPipeline = ck_tile::AQuantGemmPipelineAgBgCrCompV3; - using GemmEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem, - AccDataType, - CDataType, - ck_tile::tuple<>, - CLayout, - ck_tile::element_wise::PassThrough, - TilePartitioner::MPerBlock, - TilePartitioner::NPerBlock, - M_Warp, - N_Warp, - M_Warp_Tile, - N_Warp_Tile, - K_Warp_Tile, - transposed_warp_gemm, - ck_tile::memory_operation_enum::set>>; - using Kernel = ck_tile::QuantGemmKernel; - - auto kargs = Kernel::MakeKernelArgs(args); - - const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); - const dim3 blocks = Kernel::BlockSize(); - - if(args.k_batch != 1) - { - throw std::runtime_error("split-k is not supported yet!"); - } - - if(!Kernel::IsSupportedArgument(kargs)) - { - throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); - } - - if(s.log_level_ > 0) - { - std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' - << "shape: " << CodegenGemmShape::GetName() << '\n' - << "problem: " << CodegenPipelineProblem::GetName() << '\n' - << "pipeline: " << CodegenGemmPipeline::GetName() << '\n' - << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" - << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" - << std::endl; - } - - float ave_time = ck_tile::launch_kernel( - s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); - - return ave_time; - }; - return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); -} - -template -static constexpr inline auto is_row_major(Layout layout_) -{ - return ck_tile::bool_constant, - ck_tile::tensor_layout::gemm::RowMajor>>{}; -} - -template -float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, - ck_tile::DeviceMem& aq_m_aqk_dev_buf, - ck_tile::DeviceMem& b_k_n_dev_buf, - ck_tile::DeviceMem& c_m_n_dev_buf, - ck_tile::index_t M, - ck_tile::index_t N, - ck_tile::index_t K, - ck_tile::index_t AQK, - ck_tile::index_t stride_A, - ck_tile::index_t stride_AQ, - ck_tile::index_t stride_B, - ck_tile::index_t stride_C, - ck_tile::index_t kbatch, - int n_warmup, - int n_repeat) -{ - ck_tile::QuantGemmHostArgs args; - args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); - args.aq_ptr = aq_m_aqk_dev_buf.GetDeviceBuffer(); - args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); - args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); - args.k_batch = kbatch; - args.M = M; - args.N = N; - args.K = K; - args.QK_A = AQK; - args.stride_A = stride_A; - args.stride_B = stride_B; - args.stride_C = stride_C; - args.stride_AQ = stride_AQ; - - float ave_time = gemm_calc_aquant( - args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); - - std::size_t flop = std::size_t(2) * M * N * K; - std::size_t num_byte = sizeof(ADataType) * M * K + sizeof(AQDataType) * M * AQK + - sizeof(BDataType) * N * K + sizeof(CDataType) * M * N; - float tflops = static_cast(flop) / 1.E9 / ave_time; - float gb_per_sec = num_byte / 1.E6 / ave_time; - - std::cout << "Run Gemm kernel with M =" << M << " N =" << N << " K =" << K - << " StrideA =" << stride_A << " StrideAQ =" << stride_AQ << " StrideB =" << stride_B - << " StrideC =" << stride_C << " A_Layout =" << ALayout::name - << " B_Layout =" << BLayout::name << " C_Layout =" << CLayout::name - << " A_Type = " << DataTypeTraits::name - << " AQ_Type = " << DataTypeTraits::name - << " B_Type = " << DataTypeTraits::name - << " Acc_Type = " << DataTypeTraits::name - << " C_Type = " << DataTypeTraits::name << " : " << ave_time << " ms, " - << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; - - return ave_time; -} - -template -bool run_gemm_test_with_layouts(int argc, - char* argv[], - const ALayout a_layout = ALayout{}, - const AQLayout aq_layout = AQLayout{}, - const BLayout b_layout = BLayout{}, - [[maybe_unused]] const CLayout c_layout = CLayout{}) -{ - auto [result, arg_parser] = create_args(argc, argv); - if(!result) - return false; - - using ADataType = typename TypeConfig::ADataType; - using AQDataType = typename TypeConfig::QDataType; - using BDataType = typename TypeConfig::BDataType; - using AccDataType = typename TypeConfig::AccDataType; - using CDataType = typename TypeConfig::CDataType; - - ck_tile::index_t M = arg_parser.get_int("m"); - ck_tile::index_t N = arg_parser.get_int("n"); - ck_tile::index_t K = arg_parser.get_int("k"); - - if(K % QuantGroupSize != 0) - { - throw std::runtime_error("K must be aligned with QuantGroupSize"); - } - - ck_tile::index_t AQK = K / QuantGroupSize; - - ck_tile::index_t stride_A = arg_parser.get_int("stride_a"); - ck_tile::index_t stride_AQ = arg_parser.get_int("stride_q"); - ck_tile::index_t stride_B = arg_parser.get_int("stride_b"); - ck_tile::index_t stride_C = arg_parser.get_int("stride_c"); - - ck_tile::index_t kbatch = arg_parser.get_int("split_k"); - int n_warmup = arg_parser.get_int("warmup"); - int n_repeat = arg_parser.get_int("repeat"); - ck_tile::index_t init_method = arg_parser.get_int("init"); - - stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout)); - stride_AQ = ck_tile::get_default_stride(M, AQK, stride_AQ, is_row_major(aq_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{})); - - ck_tile::HostTensor a_m_k( - ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout))); - ck_tile::HostTensor aq_m_aqk( - ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); - ck_tile::HostTensor b_k_n( - ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(b_layout))); - ck_tile::HostTensor c_m_n_dev_result( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution fill_seed(0, 500); - - if(init_method == 0) - { - if constexpr(std::is_same_v) - { - ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}( - a_m_k); - } - else - { - ck_tile::FillUniformDistribution{-2.0f, 3.0f, fill_seed(gen)}(a_m_k); - } - ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(aq_m_aqk); - ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(b_k_n); - } - else if(init_method == 1) - { - std::cout << "Monotonic initialization is not supported." << std::endl; - return true; - } - else if(init_method == 2) - { - ck_tile::FillConstant{static_cast(0x22)}(a_m_k); - ck_tile::FillConstant{static_cast(0.5f)}(aq_m_aqk); - ck_tile::FillConstant{static_cast(0x38)}(b_k_n); - } - else - { - a_m_k.SetZero(); - aq_m_aqk.SetZero(); - b_k_n.SetZero(); - } - - ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem aq_m_aqk_dev_buf(aq_m_aqk.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - - if constexpr(std::is_same_v) - { - // Permute vector pk_i4x4 data for device implementation - ck_tile::HostTensor a_m_k_dev = a_m_k; - ck_tile::permute_vectors_i4x4_b(a_m_k_dev); - a_m_k_dev_buf.ToDevice(a_m_k_dev.data()); - } - else - { - a_m_k_dev_buf.ToDevice(a_m_k.data()); - } - aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); - b_k_n_dev_buf.ToDevice(b_k_n.data()); - c_m_n_dev_buf.SetZero(); - c_m_n_dev_result.SetZero(); - - invoke_gemm(a_m_k_dev_buf, - aq_m_aqk_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - AQK, - stride_A, - stride_AQ, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat); - - c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); - bool pass = true; - - if(arg_parser.get_int("v") == 1) - { - ck_tile::HostTensor c_m_n_host_ref( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); - c_m_n_host_ref.SetZero(); - - ck_tile::reference_gemm_quant(a_m_k, aq_m_aqk, b_k_n, c_m_n_host_ref); - const float max_accumulated_value = - *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); - const auto rtol_atol = calculate_rtol_atol( - K, kbatch, max_accumulated_value); - pass = ck_tile::check_err(c_m_n_dev_result, - c_m_n_host_ref, - "Error: Incorrect results!", - rtol_atol.at(ck_tile::number<0>{}), - rtol_atol.at(ck_tile::number<1>{})); - - if(!pass) - { - std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) - << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) - << std::endl; - } - std::cout << "CPU verification " << (pass ? "Passed!" : "Failed ...") << std::endl; - } - else if(arg_parser.get_int("v") == 2) - { - std::cout << "GPU verification is not implemented yet. Re-run with -v=1" << std::endl; - return false; - } - - return pass; -} - -template -bool run_gemm_test_prec_type(std::string a_layout, std::string b_layout, int argc, char* argv[]) -{ - using Row = ck_tile::tensor_layout::gemm::RowMajor; - using Col = ck_tile::tensor_layout::gemm::ColumnMajor; - - if constexpr(std::is_same_v || - std::is_same_v || - std::is_same_v) - { - if(a_layout == "R" && b_layout == "C") - { - return run_gemm_test_with_layouts( - argc, argv, Row{}, Row{}, Col{}, Row{}); - } - else - { - throw std::runtime_error("Unsupported memory layout for the input matrices!"); - } - } - else - { - throw std::runtime_error("Unsupported data type for A."); - } - - return true; -} - -template