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 cac08fc90f..b23815e8e8 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 @@ -103,9 +103,12 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str }(); using BaseGemmPipeline = std::decay_t; - const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile); - const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); - const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + constexpr auto K1 = GemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile) + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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_) { diff --git a/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp index 4c81718468..9202da8bd3 100644 --- a/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp @@ -180,7 +180,8 @@ struct QuantGemmHostArgs : public QuantGemmProblem const void* aq_ptr = nullptr; const void* bq_ptr = nullptr; void* c_ptr = nullptr; - index_t k_batch = 0; + // k_batch must be a positive integer; defaults to 1 (no split-K). + index_t k_batch = 1; }; struct QuantGemmKernelArgs @@ -203,10 +204,35 @@ struct QuantGemmKernelArgs index_t k_batch; }; +CK_TILE_HOST_DEVICE auto +get_splitk_batch_k_read(index_t K, index_t k_batch, index_t k_unit) noexcept -> index_t +{ + // k_batch and k_unit must be positive integers. Callers are expected to + // validate via IsSupportedArgument(); this fallback returns K so a + // misconfigured launch behaves as a no-split kernel. + if(k_batch <= 0 || k_unit <= 0) + { + return K; + } + const index_t k_t = k_batch * k_unit; + return (K + k_t - 1) / k_t * k_unit; +} + +CK_TILE_HOST_DEVICE auto +get_splitk_last_batch_k(index_t K, index_t k_batch, index_t k_read) noexcept -> index_t +{ + if(k_batch <= 0) + { + return K; + } + return K - k_read * (k_batch - 1); +} + template + QuantType QuantType_, + bool RuntimeSplitKTail_ = false> struct QuantGemmKernel { using TilePartitioner = remove_cvref_t; @@ -244,7 +270,8 @@ struct QuantGemmKernel static constexpr auto I3 = number<3>(); // BQ Tensor static constexpr auto I4 = number<4>(); // C Tensor - static constexpr auto kQuantType = QuantType_; + static constexpr auto kQuantType = QuantType_; + static constexpr bool RuntimeSplitKTail = RuntimeSplitKTail_; [[nodiscard]] CK_TILE_HOST static const std::string GetName() { @@ -386,11 +413,9 @@ struct QuantGemmKernel { constexpr auto K1 = GemmPipeline::BlockGemmShape::WarpTile::at(I2); // smallest unit of K work per block - const index_t K_t = amd_wave_read_first_lane( - kargs.k_batch * K1); // amount of K elements consumed if every split-K batch - // performs exactly one "unit" (K1) - const index_t KRead = amd_wave_read_first_lane( - (kargs.K + K_t - 1) / K_t * K1); // total k elements to be read in this batch + const index_t KRead = + amd_wave_read_first_lane(get_splitk_batch_k_read(kargs.K, kargs.k_batch, K1)); + // total k elements to be read in this batch // offset not necessarily = KRead, because B can have packed elements (e.g. fp8i4) constexpr index_t BPackedSize = ck_tile::numeric_traits>::PackedSize; @@ -412,7 +437,21 @@ struct QuantGemmKernel } else if constexpr(std::is_same_v) { - b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements); + if constexpr(PreshuffleB) + { + // Preshuffled B is laid out as [N/N_Warp_Tile, K_outer, N_Warp_Tile, K_inner] + // (see shuffle_b<>), where each "N_outer" row spans N_Warp_Tile * full_K + // linear elements. MakeBBlockWindow already builds the descriptor with + // stride [N_Warp_Tile * kargs.K, 1], so to advance the K starting position + // by k_id * KRead within row 0 we need to advance the pointer by + // (k_id * KRead) * N_Warp_Tile -- not just (k_id * KRead). + constexpr index_t N_Warp_Tile = GemmPipeline::BlockGemmShape::WarpTile::at(I1); + b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements * N_Warp_Tile); + } + else + { + b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements); + } } if(k_id < static_cast(kargs.k_batch - 1)) @@ -462,12 +501,20 @@ struct QuantGemmKernel using BQuantGroupSize = remove_cvref_t; // Compute AQ K-group offset for this split-K batch. - // AQ tensor layout is RowMajor [M, QK_A] with stride [stride_AQ, 1]. - // Advancing to column aq_group_offset means a pointer offset of aq_group_offset - // elements (column stride = 1). const index_t k_offset_aq = amd_wave_read_first_lane(k_id * KRead); - aq_group_offset = amd_wave_read_first_lane(k_offset_aq / AQuantGroupSize::kK); - aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset); + aq_group_offset = amd_wave_read_first_lane(k_offset_aq / AQuantGroupSize::kK); + if constexpr(std::is_same_v) + { + // RowMajor AQ is [M, QK_A] with stride [stride_AQ, 1]. + // Advancing to K-group column g is a pointer offset of g. + aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset); + } + else if constexpr(std::is_same_v) + { + // ColumnMajor AQ is [QK_A, M] with K-group row stride stride_AQ. + // Advancing to K-group row g is a pointer offset of g * stride_AQ. + aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset * kargs.stride_AQ); + } // Compute BQ K-group offset for this split-K batch. // BQ tensor layout is ColumnMajor [N/kN, K/kK] with stride [K/kK, 1] for @@ -1135,6 +1182,34 @@ struct QuantGemmKernel CK_TILE_HOST static bool IsSupportedArgument(const QuantGemmKernelArgs& kargs) { + // k_batch must be a positive integer. + if(kargs.k_batch <= 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("k_batch must be a positive integer (got " + + std::to_string(kargs.k_batch) + ")!"); + } + return false; + } + + // The split-K K-unit (warp-tile K dimension) must be positive too; + // it is a compile-time constant taken from the pipeline shape. + static_assert(GemmPipeline::BlockGemmShape::WarpTile::at(I2) > 0, + "Pipeline warp-tile K dimension (k_unit) must be positive."); + + // ABQuantGrouped does not currently support RowMajor BQ layout: the + // BQ tensor view, tile window, and split-K offset code are all + // written for ColumnMajor BQ. The deeper static_asserts in + // MakeBQBlockWindow enforce this at instantiation time; surface it + // here at the host-arg entry point too so the limitation is visible + // before the first device-side instantiation. + static_assert(!(kQuantType == QuantType::ABQuantGrouped && + std::is_same_v), + "ABQuantGrouped does not currently support RowMajor BQ layout. " + "Use ColumnMajor BQ (or extend MakeBQBlockWindow and the split-K " + "BQ offset path to handle RowMajor BQ)."); + // Split-K is supported for BQuantGrouped (without preshuffle) and // ABQuantGrouped (without APreshuffleQuant) modes. if(kargs.k_batch != 1) @@ -1158,12 +1233,22 @@ struct QuantGemmKernel } else { - constexpr auto K1 = GemmPipeline::BlockGemmShape::WarpTile::at(I2); - const index_t K_t = kargs.k_batch * K1; - const index_t KRead = (kargs.K + K_t - 1) / K_t * K1; // per-batch K read size + constexpr auto K1 = GemmPipeline::BlockGemmShape::WarpTile::at(I2); + const index_t KRead = + get_splitk_batch_k_read(kargs.K, kargs.k_batch, K1); // per-batch K read size + const index_t KLast = get_splitk_last_batch_k(kargs.K, kargs.k_batch, KRead); constexpr index_t BPackedSize = ck_tile::numeric_traits>::PackedSize; + if(KLast <= 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("Split-K configuration produces an empty final K batch!"); + } + return false; + } + // Constraint 1: KRead must align with B packing requirements. // For packed data types, multiple K elements are stored in each storage unit. // Split-K advances the B pointer by (KRead / BPackedSize) storage units per batch. @@ -1223,8 +1308,7 @@ struct QuantGemmKernel // (i.e. per_batch_num_loop == 1) the prefetch would read the tile // belonging to the next split-K batch, producing incorrect results. { - const index_t per_batch_num_loop = - KRead / static_cast(TilePartitioner::KPerBlock); + const index_t per_batch_num_loop = TilePartitioner::GetLoopNum(KRead); if(per_batch_num_loop < 2) { if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) @@ -1240,6 +1324,33 @@ struct QuantGemmKernel return false; } } + + // Host-side fixed tail selection is only valid when all split-K batches have + // the same hot-loop/tail classification. Earlier batches use KRead; the final + // batch may be shorter due to split rounding. + { + const index_t first_num_loop = TilePartitioner::GetLoopNum(KRead); + const index_t last_num_loop = TilePartitioner::GetLoopNum(KLast); + const bool first_hot_loop = GemmPipeline::BlockHasHotloop(first_num_loop); + const bool last_hot_loop = GemmPipeline::BlockHasHotloop(last_num_loop); + const auto first_tail = GemmPipeline::GetBlockLoopTailNum(first_num_loop); + const auto last_tail = GemmPipeline::GetBlockLoopTailNum(last_num_loop); + + if constexpr(!RuntimeSplitKTail) + { + if(first_hot_loop != last_hot_loop || first_tail != last_tail) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Split-K batches require different hot-loop/tail handling. " + "Use a K/k_batch combination that gives matching pipeline " + "tails or enable runtime split-K tail dispatch."); + } + return false; + } + } + } } } @@ -1383,6 +1494,78 @@ struct QuantGemmKernel return true; } + template + CK_TILE_DEVICE static auto CallBQuantGemmPipeline(const ADramBlockWindow& a_block_window, + const BDramBlockWindow& b_block_window, + const BQDramBlockWindow& bq_block_window, + const index_t num_loop, + void* smem_ptr, + const index_t n) + { + if constexpr(RuntimeSplitKTail) + { + static_assert(!PreshuffleB, + "RuntimeSplitKTail is not implemented for preshuffle-B BQuant " + "pipelines."); + const bool has_hot_loop = GemmPipeline::BlockHasHotloop(num_loop); + const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop); + return GemmPipeline{}(a_block_window, + b_block_window, + bq_block_window, + num_loop, + has_hot_loop, + tail_num, + smem_ptr, + n); + } + else + { + return GemmPipeline{}( + a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr, n); + } + } + + template + CK_TILE_DEVICE static auto CallABQuantGemmPipeline(const ADramBlockWindow& a_block_window, + const BDramBlockWindow& b_block_window, + const AQDramBlockWindow& aq_block_window, + const BQDramBlockWindow& bq_block_window, + const index_t num_loop, + void* smem_ptr, + const index_t m, + const index_t n) + { + if constexpr(RuntimeSplitKTail) + { + const bool has_hot_loop = GemmPipeline::BlockHasHotloop(num_loop); + const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop); + return GemmPipeline{}(a_block_window, + b_block_window, + aq_block_window, + bq_block_window, + num_loop, + has_hot_loop, + tail_num, + smem_ptr, + m, + n); + } + else + { + return GemmPipeline{}(a_block_window, + b_block_window, + aq_block_window, + bq_block_window, + num_loop, + smem_ptr, + m, + n); + } + } + /** * @brief Runs single GEMM problem cooperatively by whole workgroup. * @@ -1425,7 +1608,6 @@ struct QuantGemmKernel const index_t num_loop = amd_wave_read_first_lane(TilePartitioner::GetLoopNum(splitk_batch_offset.splitted_k)); - // Run GEMM cooperatively by whole workgroup. const auto& c_block_tile = [&]() { if constexpr(kQuantType == QuantType::AQuantGrouped) @@ -1445,7 +1627,7 @@ struct QuantGemmKernel { n = kargs.N; } - return GemmPipeline{}( + return CallBQuantGemmPipeline( a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr, n); } else if constexpr(kQuantType == QuantType::ABQuantGrouped) @@ -1457,14 +1639,14 @@ struct QuantGemmKernel // m = kargs.M; n = kargs.N; } - return GemmPipeline{}(a_block_window, - b_block_window, - aq_block_window, - bq_block_window, - num_loop, - smem_ptr, - m, - n); + return CallABQuantGemmPipeline(a_block_window, + b_block_window, + aq_block_window, + bq_block_window, + num_loop, + smem_ptr, + m, + n); } else if constexpr(kQuantType == QuantType::RowColQuant || kQuantType == QuantType::TensorQuant) @@ -1557,7 +1739,6 @@ struct QuantGemmKernel // allocate LDS __shared__ char smem_ptr[GetSmemSize()]; - assert(kargs.k_batch == 1); RunGemm( a_ptr, b_ptr, aq_ptr, bq_ptr, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n); } diff --git a/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp index 46284c62ae..be9d94d8b8 100644 --- a/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp @@ -314,12 +314,18 @@ struct QuantGroupedGemmKernel const typename Base::SplitKBatchOffset splitk_batch_offset(kargs, block_idx_z); - // options - const ADataType* a_ptr = static_cast(kargs.a_ptr); - const BDataType* b_ptr = static_cast(kargs.b_ptr); - const AQDataType* aq_ptr = static_cast(kargs.aq_ptr); - const BQDataType* bq_ptr = static_cast(kargs.bq_ptr); - CDataType* c_ptr = static_cast(kargs.c_ptr); + // Apply split-K per-batch offsets to every input pointer, mirroring the + // non-grouped QuantGemmKernel::Run_ path. All offsets are 0 when + // k_batch == 1, so this is a no-op for non-split-K launches. + const ADataType* a_ptr = + static_cast(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset; + const BDataType* b_ptr = + static_cast(kargs.b_ptr) + splitk_batch_offset.b_k_split_offset; + const AQDataType* aq_ptr = + static_cast(kargs.aq_ptr) + splitk_batch_offset.aq_k_split_offset; + const BQDataType* bq_ptr = + static_cast(kargs.bq_ptr) + splitk_batch_offset.bq_k_split_offset; + CDataType* c_ptr = static_cast(kargs.c_ptr); // allocate LDS __shared__ char smem_ptr[GetSmemSize()]; @@ -454,8 +460,11 @@ struct QuantGroupedGemmKernel Base::MakeABlockWindow(a_ptr, kargs, splitk_batch_offset.splitted_k, block_idx_m); const auto& b_block_window = Base::MakeBBlockWindow(b_ptr, kargs, splitk_batch_offset.splitted_k, block_idx_n); - const auto& aq_block_window = - Base::MakeAQBlockWindow(aq_ptr, kargs, block_idx_m, block_idx_n); + // Pass aq_group_offset so the AQ tensor view dimension reflects the + // remaining K-groups from this split-K batch's offset position + // (mirrors QuantGemmKernel::RunGemm). + const auto& aq_block_window = Base::MakeAQBlockWindow( + aq_ptr, kargs, block_idx_m, block_idx_n, splitk_batch_offset.aq_group_offset); const auto& bq_block_window = Base::MakeBQBlockWindow( bq_ptr, kargs, splitk_batch_offset.bq_group_offset, block_idx_m, block_idx_n); diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_v2.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_v2.hpp index f6e7f14bc3..dcbfd7a172 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_v2.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_v2.hpp @@ -637,8 +637,10 @@ struct WPABQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRe const AQDramBlockWindowTmp& aq_dram_block_window_tmp, const BQDramBlockWindowTmp& bq_dram_block_window_tmp, index_t num_loop, + bool has_hot_loop, TailNumber tail_number, void* p_smem, + index_t m = 0, index_t n = 0) const { const auto RunPipeline = [&](auto bool_val, auto tail_num_) { @@ -650,11 +652,14 @@ struct WPABQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRe b_flat_dram_block_window_tmp, aq_dram_block_window_tmp, bq_dram_block_window_tmp, - n, // dummy value, won't be used + // The preshuffle-B ABQuant pipeline currently ignores m and n; keep this + // runtime-tail wrapper aligned with the generic ABQuant pipeline signature. + m, + n, num_loop, p_smem); }; - return Base::TailHandler(RunPipeline, true, tail_number); + return Base::TailHandler(RunPipeline, has_hot_loop, tail_number); } }; } // namespace ck_tile diff --git a/test/ck_tile/gemm_block_scale/CMakeLists.txt b/test/ck_tile/gemm_block_scale/CMakeLists.txt index 676f01e355..ab6708793a 100644 --- a/test/ck_tile/gemm_block_scale/CMakeLists.txt +++ b/test/ck_tile/gemm_block_scale/CMakeLists.txt @@ -92,6 +92,11 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12") ) target_compile_options(test_tile_gemm_quant_abquant_splitk_prefill PRIVATE ${TEST_GEMM_COMPILE_OPTIONS}) + add_gtest_executable(test_tile_gemm_quant_abquant_splitk_preshuffleB + test_gemm_quant_abquant_splitk_preshuffleB.cpp + ) + target_compile_options(test_tile_gemm_quant_abquant_splitk_preshuffleB PRIVATE ${TEST_GEMM_COMPILE_OPTIONS}) + add_gtest_executable(test_tile_gemm_quant_abquant_a4w4_base test_gemm_quant_abquant_a4w4_base.cpp ) @@ -278,6 +283,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12") # ABQuant split-K tests test_tile_gemm_quant_abquant_splitk_decode test_tile_gemm_quant_abquant_splitk_prefill + test_tile_gemm_quant_abquant_splitk_preshuffleB # BQuant tests test_tile_gemm_quant_bquant_1d_128 test_tile_gemm_quant_bquant_1d_64 diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_decode.cpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_decode.cpp index 7f8fb70f99..42cb5a3597 100644 --- a/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_decode.cpp +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_decode.cpp @@ -110,3 +110,46 @@ TYPED_TEST(TestCkTileGemmABQuant, SplitK8_LargeK_LargeMN) // K=4096, larger M and N this->run_test_with_validation(48, 192, 4096, 8); } + +// Test one padded-K split shape whose earlier split-K batches and final batch would need +// different compile-time pipeline tail handling. Fixed host-side tail selection rejects it. +// +// K=3328, k_batch=9, GemmConfigPadding (K_Tile=256): +// K_Warp_Tile is arch-dependent (gfx94/95: 32 or 64, gfx12 WMMA: 16); for all of these +// ceil(3328 / (9 * K_Warp_Tile)) * K_Warp_Tile = 384, so KRead is always a multiple of +// BQuantGroupSize::kK = AQuantGroupSize::kK = 128 (Constraints 2/3 of IsSupportedArgument). +// KLast = 3328 - 8*384 = 256. +// num_loop_first = ceil(384/256) = 2 (hot_loop=false, tail=Even) +// num_loop_last = ceil(256/256) = 1 (hot_loop=false, tail=Odd) +// Mismatched tail, so the fixed host-side dispatch rejects; the runtime-tail dispatch path +// (RuntimeSplitKTail=true) accepts and dispatches per-batch. +using ABQuantSplitKRejectTypes = ::testing::Types>; + +template +class TestCkTileGemmABQuantSplitKReject : public TestCkTileGemmABQuant +{ +}; + +TYPED_TEST_SUITE(TestCkTileGemmABQuantSplitKReject, ABQuantSplitKRejectTypes); + +TYPED_TEST(TestCkTileGemmABQuantSplitKReject, RejectsMismatchedTailSplitK) +{ + EXPECT_THROW(this->run_test_with_validation(32, 128, 3328, 9), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmABQuantSplitKReject, RuntimeTailAllowsMismatchedTailSplitK) +{ + this->run_test_with_validation(32, 128, 3328, 9, 0, true /* allow_runtime_splitk_tail */); +} diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_preshuffleB.cpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_preshuffleB.cpp new file mode 100644 index 0000000000..9091b6acb3 --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_preshuffleB.cpp @@ -0,0 +1,94 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_quant_common.hpp" + +using GroupSize1x1x128 = ck_tile::QuantGroupShape>; +using GroupSize1x128x128 = ck_tile::QuantGroupShape>; + +// ABQuant split-K with B-preshuffle pipeline (WPABQuantBPipelineAgBgCrV2). +// Exercises both the regular (uniform-split) and the runtime-tail (uneven-split) +// dispatch paths, mirroring the non-preshuffle reject/accept tests in +// test_gemm_quant_abquant_splitk_decode.cpp. +// +// Tuple format: + +// ===================================================================================== +// Uniform-split tests: every split-K batch has the same per-batch num_loop and tail +// classification, so the regular (non-runtime-tail) dispatch path applies. +// ===================================================================================== + +// clang-format off +using ABQuantSplitKPreshuffleBUniformTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +TYPED_TEST_SUITE(TestCkTileGemmABQuant, ABQuantSplitKPreshuffleBUniformTypes); + +// GemmConfigPreshuffleBPrefill: M_Tile=128, N_Tile=128, K_Tile=128, K_Warp_Tile=64. +// For uniform splits we want every batch to have the same num_loop classification: +// pick K such that KRead == KLast and num_loop_per_batch >= 2. + +TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK2_K1024) +{ + // K=1024, k_batch=2 -> KRead=KLast=512, num_loop=4 per batch (Even tail). + this->run_test_with_validation(128, 128, 1024, 2); +} + +TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK4_K2048) +{ + // K=2048, k_batch=4 -> KRead=KLast=512, num_loop=4 per batch (Even tail). + this->run_test_with_validation(128, 128, 2048, 4); +} + +TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK2_LargeK_LargeN) +{ + // K=2048, larger N (multiple of N_Tile=128). + this->run_test_with_validation(128, 256, 2048, 2); +} + +// ===================================================================================== +// Runtime-tail tests: K and k_batch chosen so the first split-K batch and the final +// (shorter) batch land in different (hot-loop, tail) classifications. The default +// host-side fixed tail dispatch must reject this; the runtime-tail dispatch path +// (RuntimeSplitKTail=true) must accept it. Uses the padded preshuffle config so +// uneven K passes the kPadK=false divisibility check (mirrors the non-preshuffle +// decode test against GemmConfigPadding). +// ===================================================================================== + +// clang-format off +using ABQuantSplitKPreshuffleBRuntimeTailTypes = ::testing::Types< + std::tuple +>; +// clang-format on + +template +class TestCkTileGemmABQuantSplitKPreshuffleBReject : public TestCkTileGemmABQuant +{ +}; + +TYPED_TEST_SUITE(TestCkTileGemmABQuantSplitKPreshuffleBReject, + ABQuantSplitKPreshuffleBRuntimeTailTypes); + +// K=3328, k_batch=9 with K_Tile=128: +// K_Warp_Tile is arch-dependent (gfx94/95: 64 for 8-bit, gfx12 WMMA: 16); for both +// ceil(3328 / (9 * K_Warp_Tile)) * K_Warp_Tile = 384, so KRead is always a multiple of +// BQuantGroupSize::kK = AQuantGroupSize::kK = 128 (Constraints 2/3 of IsSupportedArgument). +// KLast = 3328 - 8*384 = 256. +// num_loop_first = ceil(384/128) = 3 (hot_loop=true, tail=Odd) +// num_loop_last = ceil(256/128) = 2 (hot_loop=false, tail=Even) +// Both hot_loop and tail differ, so the fixed host-side dispatch rejects; the runtime-tail +// dispatch path (RuntimeSplitKTail=true) accepts and dispatches per-batch. + +TYPED_TEST(TestCkTileGemmABQuantSplitKPreshuffleBReject, RejectsMismatchedTailSplitK) +{ + EXPECT_THROW(this->run_test_with_validation(128, 128, 3328, 9), std::runtime_error); +} + +TYPED_TEST(TestCkTileGemmABQuantSplitKPreshuffleBReject, RuntimeTailAllowsMismatchedTailSplitK) +{ + this->run_test_with_validation(128, 128, 3328, 9, 0, true /* allow_runtime_splitk_tail */); +} 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 index d3973e3080..c3649645d1 100644 --- a/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp @@ -53,6 +53,31 @@ struct SafeTupleElement using SafeTupleElement_t = typename SafeTupleElement::type; +namespace test_gemm_quant_base_detail { +// TODO: replace with C++20 requires later. +// C++17 detection idiom: true when +// T::run_quant_gemm_impl(QuantGemmHostArgs, stream_config, bool) +// is a well-formed expression. +template +struct has_run_quant_gemm_impl_splitk : std::false_type +{ +}; + +template +struct has_run_quant_gemm_impl_splitk< + T, + Shape, + Partitioner, + Traits, + std::void_t< + decltype(std::declval()->template run_quant_gemm_impl( + std::declval(), + std::declval(), + std::declval()))>> : std::true_type +{ +}; +} // namespace test_gemm_quant_base_detail + // Base class for common quant gemm functionality template class TestCkTileGemmQuantBase : public ::testing::Test @@ -114,7 +139,9 @@ class TestCkTileGemmQuantBase : public ::testing::Test 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) + void invoke_quant_gemm(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s, + bool allow_runtime_splitk_tail = false) { // WP pipeline requires per-thread tile size aligned to Problem::VectorLoadSize. // static_assert((WG::kM * WG::kK * sizeof(ADataType) * MIterPerWarp / WaveSize) % @@ -149,9 +176,24 @@ class TestCkTileGemmQuantBase : public ::testing::Test VectorSize>; // Let the derived class create the appropriate pipeline and epilogue - static_cast(this) - ->template run_quant_gemm_impl( - args, s); + auto* derived = static_cast(this); + if constexpr(test_gemm_quant_base_detail::has_run_quant_gemm_impl_splitk< + Derived, + CodegenGemmShape, + TilePartitioner, + CodegenGemmTraits>::value) + { + derived->template run_quant_gemm_impl( + args, s, allow_runtime_splitk_tail); + } + else + { + derived->template run_quant_gemm_impl(args, s); + } } void RunTest(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) 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 index 42d850c6bb..1529faf86a 100644 --- a/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp @@ -3,6 +3,8 @@ #pragma once +#include + #include "test_gemm_quant_base.hpp" #include "ck_tile/host/permute_pk_int4.hpp" #include "ck_tile/host/tensor_shuffle_utils.hpp" @@ -145,6 +147,12 @@ struct GemmConfigPreshuffleBPrefillTransposeC : public GemmConfigPreshuffleBPref static constexpr bool TransposeC = true; }; +struct GemmConfigPreshuffleBPrefillPadded : public GemmConfigPreshuffleBPrefill +{ + static constexpr bool kPadN = true; + static constexpr bool kPadK = true; +}; + struct GemmConfigPreshuffleQuantPrefill : public GemmConfigPrefill { static constexpr bool BPreshuffleQuant = true; @@ -375,9 +383,12 @@ class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase; - 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); + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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_) { @@ -594,16 +605,19 @@ class TestCkTileGemmAQuantMem 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::BaseGemmPipelineAgBgCrMem; - 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); + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem; + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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; @@ -888,9 +902,12 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase, ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2>; - 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); + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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_) { @@ -1022,8 +1039,9 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBaseis_row_major(ALayout{})); @@ -1167,7 +1185,7 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBaseinvoke_quant_gemm(args, stream_config); + this->invoke_quant_gemm(args, stream_config, allow_runtime_splitk_tail); // Validation using reference implementation ck_tile::HostTensor c_m_n_host_ref( @@ -1216,11 +1234,16 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, - const ck_tile::stream_config& s) + const ck_tile::stream_config& s, + bool allow_runtime_splitk_tail) { static_assert(std::is_same_v); @@ -1258,8 +1281,10 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase; + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); const ck_tile::index_t K_split = - ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile); + args.k_batch == 1 ? ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile) + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); 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); @@ -1330,23 +1355,37 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase>>; - using Kernel = ck_tile::QuantGemmKernel; + // TODO: Replace with templated lambda when C++20 is available + auto LaunchKernel = [&](auto RuntimeSplitKTailTag) { + constexpr bool RuntimeSplitKTail = decltype(RuntimeSplitKTailTag)::value; + 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(); + 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)) + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for ABQuant kernel"); + } + using k_attr_t = ck_tile::kernel_attr; + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + if(allow_runtime_splitk_tail) { - throw std::runtime_error("Arguments not supported for ABQuant kernel"); + LaunchKernel(std::true_type{}); + } + else + { + LaunchKernel(std::false_type{}); } - using k_attr_t = ck_tile::kernel_attr; - ck_tile::launch_kernel(s, - ck_tile::make_kernel( - Kernel{}, grids, blocks, 0, kargs)); }; return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); @@ -1510,9 +1549,12 @@ class TestCkTileGemmRowColQuant 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); + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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_) { @@ -1725,9 +1767,12 @@ class TestCkTileGemmTensorQuant 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); + constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = + args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile + : ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1); + 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_) { diff --git a/test/ck_tile/grouped_gemm_quant/CMakeLists.txt b/test/ck_tile/grouped_gemm_quant/CMakeLists.txt index 04db7b0c4b..b52b79b0f9 100644 --- a/test/ck_tile/grouped_gemm_quant/CMakeLists.txt +++ b/test/ck_tile/grouped_gemm_quant/CMakeLists.txt @@ -23,6 +23,16 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12") add_gtest_executable(test_ck_tile_grouped_gemm_quant_bquant_preshuffleb test_grouped_gemm_quant_bquant_preshuffleb.cpp) target_compile_options(test_ck_tile_grouped_gemm_quant_bquant_preshuffleb PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + # Smoke test for grouped ABQuant + split-K. Verifies that + # QuantGroupedGemmKernel<..., ABQuantGrouped, ...> instantiates, that + # IsSupportedArgument behaves correctly for valid / invalid k_batch, and + # runs one minimal end-to-end correctness check (single group, k_batch=2) + # to cover the grouped-specific per-batch pointer offsetting and AQ + # group offset wiring (broader code paths are tested under + # gemm_block_scale/test_gemm_quant_abquant_splitk_*). + add_gtest_executable(test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke test_grouped_gemm_quant_abquant_splitk_smoke.cpp) + target_compile_options(test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + # Collect all test targets for umbrella label set(CK_TILE_GROUPED_GEMM_QUANT_TEST_TARGETS test_ck_tile_grouped_gemm_quant_rowcol @@ -30,6 +40,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12") test_ck_tile_grouped_gemm_quant_aquant test_ck_tile_grouped_gemm_quant_bquant test_ck_tile_grouped_gemm_quant_bquant_preshuffleb + test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke ) # Label all ck_tile grouped_gemm_quant tests with CK_TILE_GROUPED_GEMM_QUANT_TESTS for selective execution diff --git a/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_quant_abquant_splitk_smoke.cpp b/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_quant_abquant_splitk_smoke.cpp new file mode 100644 index 0000000000..d4e5cdb12f --- /dev/null +++ b/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_quant_abquant_splitk_smoke.cpp @@ -0,0 +1,381 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// Minimal smoke test for QuantGroupedGemmKernel<..., ABQuantGrouped, ...> +// with split-K. The main ABQuant + split-K code paths (uniform splits, +// runtime-tail dispatch, BPreshuffle, etc.) are exercised in +// test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_*.cpp using +// the non-grouped kernel; the grouped kernel reuses the same Base::RunGemm. +// What we cover here is grouped-specific: +// +// 1. Compile-time instantiation of the grouped kernel for ABQuantGrouped. +// If any inner static_assert (e.g. RowMajor BQ + ABQuant) fires this +// file won't compile. +// 2. Host-side IsSupportedArgument acceptance for valid k_batch and +// rejection for k_batch <= 0. +// 3. A single end-to-end correctness run with k_batch == 2 on a single +// group. This exercises QuantGroupedGemmKernel::Run's per-batch +// pointer offsetting (a/b/aq/bq) and the aq_group_offset wired into +// MakeAQBlockWindow -- both of which were latent bugs before this +// change because every existing grouped test launched with k_batch=1. + +#include "ck_tile/host.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm_quant.hpp" + +#include + +namespace { + +using ALayout = ck_tile::tensor_layout::gemm::RowMajor; +using BLayout = ck_tile::tensor_layout::gemm::ColumnMajor; +using CLayout = ck_tile::tensor_layout::gemm::RowMajor; +using AQLayout = ck_tile::tensor_layout::gemm::RowMajor; +using BQLayout = ck_tile::tensor_layout::gemm::ColumnMajor; +using ADataType = ck_tile::fp8_t; +using BDataType = ck_tile::fp8_t; +using AccDataType = float; +using CDataType = ck_tile::half_t; +using QDataType = float; + +using ComputeDataType = ADataType; +using AQuantGroupSize = ck_tile::QuantGroupShape>; +using BQuantGroupSize = ck_tile::QuantGroupShape>; + +constexpr ck_tile::index_t M_Tile = 128; +constexpr ck_tile::index_t N_Tile = 128; +constexpr ck_tile::index_t K_Tile = 128; +constexpr ck_tile::index_t M_Warp = 1; +constexpr ck_tile::index_t N_Warp = 4; +constexpr ck_tile::index_t K_Warp = 1; +constexpr ck_tile::index_t M_Warp_Tile = 16; +constexpr ck_tile::index_t N_Warp_Tile = 16; +#if CK_TILE_USE_WMMA +constexpr ck_tile::index_t K_Warp_Tile = 16; +#else +constexpr ck_tile::index_t K_Warp_Tile = 64; +#endif +constexpr bool kPadM = false; +constexpr bool kPadN = false; +constexpr bool kPadK = false; +constexpr bool TransposeC = true; +constexpr ck_tile::QuantType QuantMode = ck_tile::QuantType::ABQuantGrouped; + +using GemmShape = ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + +using TilePartitioner = ck_tile::GemmTile1DPartitioner; + +using GemmTraits = ck_tile::TileGemmQuantTraits; + +// PipelineProblem template used to drive BaseGemmPipeline / TailHandler +// dispatch. The HasHotLoop / TailNumber template parameters are fixed at +// compile time per-instantiation; TailHandler picks the right one at +// runtime based on (K, k_batch). +template +using AbquantPipelineProblem = + ck_tile::GemmABQuantPipelineProblem; + +using BasePipelineProblem = ck_tile::GemmPipelineProblemBase; +using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + +// Compile-time instantiation check: a representative set of types must +// build cleanly (this is the minimal "smoke" portion of the test). +using SmokePipeline = ck_tile::ABQuantGemmPipelineAgBgCrCompV3< + AbquantPipelineProblem>; +using SmokeEpilogue = + ck_tile::CShuffleEpilogue, + 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, + TransposeC>>; + +using SmokeGroupedKernel = + ck_tile::QuantGroupedGemmKernel; +using SmokeBaseKernel = typename SmokeGroupedKernel::Base; + +static_assert(sizeof(SmokeGroupedKernel) > 0, "QuantGroupedGemmKernel must instantiate"); +static_assert(sizeof(SmokeBaseKernel) > 0, "QuantGemmKernel base must instantiate"); + +ck_tile::QuantGemmKernelArgs MakeKargsForValidation(ck_tile::index_t k_batch) +{ + constexpr ck_tile::index_t M = 128; + constexpr ck_tile::index_t N = 128; + constexpr ck_tile::index_t K = 1024; + + ck_tile::QuantGemmKernelArgs kargs{}; + kargs.a_ptr = nullptr; + kargs.b_ptr = nullptr; + kargs.aq_ptr = nullptr; + kargs.bq_ptr = nullptr; + kargs.c_ptr = nullptr; + kargs.M = M; + kargs.N = N; + kargs.K = K; + kargs.QK_A = ck_tile::integer_divide_ceil(K, AQuantGroupSize::kK); + kargs.QK_B = ck_tile::integer_divide_ceil(K, BQuantGroupSize::kK); + kargs.stride_A = K; + kargs.stride_B = K; + kargs.stride_C = N; + kargs.stride_AQ = kargs.QK_A; + kargs.stride_BQ = N; + kargs.k_batch = k_batch; + return kargs; +} + +// End-to-end runner. Builds host tensors, fills them with deterministic +// random data, runs the grouped kernel for a single group, and validates +// against the host reference. Kept intentionally narrow: one group, fixed +// layouts/types, no preshuffle. The wider parameter space is covered by +// the non-grouped ABQuant tests. +bool RunSingleGroupABQuantSplitK(ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + ck_tile::index_t k_batch) +{ + auto is_row_major = [](auto layout) { + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; + }; + + const ck_tile::index_t stride_A = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{})); + const ck_tile::index_t stride_B = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{})); + const ck_tile::index_t stride_C = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{})); + + const ck_tile::index_t AQK = ck_tile::integer_divide_ceil(K, AQuantGroupSize::kK); + const ck_tile::index_t BQN = ck_tile::integer_divide_ceil(N, BQuantGroupSize::kN); + const ck_tile::index_t BQK = ck_tile::integer_divide_ceil(K, BQuantGroupSize::kK); + const ck_tile::index_t stride_AQ = + ck_tile::get_default_stride(M, AQK, 0, is_row_major(AQLayout{})); + const ck_tile::index_t stride_BQ = + ck_tile::get_default_stride(BQK, BQN, 0, is_row_major(BQLayout{})); + + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(BLayout{}))); + ck_tile::HostTensor aq_m_aqk( + ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(AQLayout{}))); + ck_tile::HostTensor bq_bqk_bqn( + ck_tile::host_tensor_descriptor(BQK, BQN, stride_BQ, is_row_major(BQLayout{}))); + + 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); + ck_tile::FillUniformDistribution{-2.0f, 2.0f}(bq_bqk_bqn); + + 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 aq_m_aqk_dev_buf(aq_m_aqk.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem bq_bqk_bqn_dev_buf(bq_bqk_bqn.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); + bq_bqk_bqn_dev_buf.ToDevice(bq_bqk_bqn.data()); + + if(k_batch > 1) + { + c_m_n_dev_buf.SetZero(); + } + + std::vector gemm_descs; + gemm_descs.emplace_back(a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + c_m_n_dev_buf.GetDeviceBuffer(), + aq_m_aqk_dev_buf.GetDeviceBuffer(), + bq_bqk_bqn_dev_buf.GetDeviceBuffer(), + k_batch, + M, + N, + K, + AQK, + BQK, + stride_A, + stride_B, + stride_C, + stride_AQ, + stride_BQ); + + // Workspace holds the per-group QuantGemmTransKernelArg vector, copied + // to device before the launch. + ck_tile::DeviceMem gemm_workspace(gemm_descs.size() * sizeof(ck_tile::QuantGemmTransKernelArg)); + void* kargs_ptr = gemm_workspace.GetDeviceBuffer(); + + // Drive TailHandler dispatch with split-K-aware K_split (mirrors + // run_quant_gemm_impl in the non-grouped fixture). + constexpr auto K1 = GemmShape::WarpTile::at(ck_tile::number<2>{}); + const ck_tile::index_t K_split = (k_batch == 1) + ? ck_tile::integer_least_multiple(K, K_Tile) + : ck_tile::get_splitk_batch_k_read(K, k_batch, K1); + 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); + + ck_tile::stream_config s{}; + + 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 = AbquantPipelineProblem; + using GemmPipeline = ck_tile::ABQuantGemmPipelineAgBgCrCompV3; + 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, + TransposeC>>; + using Kernel = + ck_tile::QuantGroupedGemmKernel; + + auto kargs = Kernel::MakeKargs(gemm_descs); + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Grouped ABQuant SplitK args not supported"); + } + + const dim3 grids = Kernel::GridSize(gemm_descs); + const dim3 blocks = Kernel::BlockSize(); + + HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr, + kargs.data(), + kargs.size() * sizeof(ck_tile::QuantGemmTransKernelArg), + hipMemcpyHostToDevice, + s.stream_id_)); + + ck_tile::launch_kernel( + s, + ck_tile::make_kernel<1>(Kernel{}, + grids, + blocks, + 0, + ck_tile::cast_pointer_to_constant_address_space(kargs_ptr), + static_cast(gemm_descs.size()))); + }; + + BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + + 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_abquant( + a_m_k, aq_m_aqk, b_k_n, bq_bqk_bqn, c_m_n_host_ref); + + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + 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 = + std::max(ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, k_batch)), + ck_tile::get_relative_threshold(k_batch)); + const auto atol = + std::max(ck_tile::get_absolute_threshold( + max_accumulated_value / k_batch, ck_tile::integer_divide_ceil(K, k_batch)), + ck_tile::get_absolute_threshold( + max_accumulated_value, k_batch)); + + return ck_tile::check_err( + c_m_n_dev_result, c_m_n_host_ref, "Grouped ABQuant SplitK mismatch", rtol, atol); +} + +} // namespace + +TEST(GroupedABQuantSplitKSmoke, AcceptsKBatch1) +{ + EXPECT_TRUE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/1))); +} + +TEST(GroupedABQuantSplitKSmoke, AcceptsKBatch2) +{ + EXPECT_TRUE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/2))); +} + +TEST(GroupedABQuantSplitKSmoke, RejectsKBatchZero) +{ + EXPECT_FALSE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/0))); +} + +TEST(GroupedABQuantSplitKSmoke, RejectsKBatchNegative) +{ + EXPECT_FALSE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/-1))); +} + +// End-to-end correctness for the grouped ABQuant + split-K path on a +// single group with k_batch=2. Catches regressions in: +// - QuantGroupedGemmKernel::Run per-batch a/b/aq/bq pointer offsetting, +// - aq_group_offset wiring through MakeAQBlockWindow. +// Both were latent (existing grouped tests only used k_batch=1). +TEST(GroupedABQuantSplitKSmoke, EndToEnd_SingleGroup_KBatch2) +{ + EXPECT_TRUE(RunSingleGroupABQuantSplitK(/*M=*/128, + /*N=*/128, + /*K=*/512, + /*k_batch=*/2)); +}