diff --git a/include/ck_tile/ops/gemm_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 index c010affd61..49599ea3de 100644 --- a/include/ck_tile/ops/gemm_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 @@ -76,6 +76,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; + static constexpr index_t NQPerBlock = NPerBlock / QuantGroupSize::kN; static constexpr index_t BQPerBlock = KPerBlock / QuantGroupSize::kK; static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); @@ -340,23 +341,32 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase } }); - // Need to multiply bquant with accumulated C - // - // The accumulated C tile has the standard distribution. For example - // lane 0 holds elements [0,0], [1,0], [2,0], [3,0], [8,0], [9,0], - // [10,0], [11,0], [16,0], [17,0], [18,0], [19,0], [24,0], [25,0], - // [26,0], [27,0]. - // - // These elements are in different rows, need to get the scale value - // for the corresponding row. - // Based on bquant's tile distribution, it can be inferred which - // lane holds the relevant scale. For example, the scales corresponding - // to the 16 elements held by lane 0 are held by lanes 0, 1, 2, 3, 8, 9, - // 10, 11, 16, 17, 18, 19, 24, 25, 26, 27 respectively. - // - // These scales can be obtained using __builtin_amdgcn_ds_bpermute. - - constexpr index_t reg_offset = nIter * Traits::BQPerBlock + kQScale; + // Multiply bquant with accumulated C + const index_t reg_offset = [&]() { + if constexpr(Traits::NQPerBlock == Traits::NPerBlock) + { + // Each row of B has a separate scale, each thread has its own single element + // of the scale matrix for the current nIter/kQScale + return nIter * Traits::BQPerBlock + kQScale; + } + else + { + // FIXME: temporarily the tile distribution replicates all block's scales + // to all threads - need to figure out the index manually here + // from nIter and warp id + const index_t n_idx_of_warp = nIter * WarpGemm::kN * NWarp + + get_warp_id() * WarpGemm::kN; + const index_t row_index = n_idx_of_warp / Traits::QuantGroupSize::kN; + if (threadIdx.x == 0) + { + printf("n_idx_of_warp: %d, row_index: %d, kQScale: %d\n", + n_idx_of_warp, + row_index, + kQScale.value); + } + return row_index * Traits::BQPerBlock + kQScale; + } + }(); constexpr auto tbuf_offset = number; return TileEncodingPattern::make_2d_static_tile_distribution(); diff --git a/include/ck_tile/ops/gemm_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 index 57e0cc3e1b..d2c0d5ced8 100644 --- a/include/ck_tile/ops/gemm_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 @@ -94,7 +94,7 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseBQuantGemmPipelineAgBgCrCompV using QuantGroupSize = remove_cvref_t; static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!"); - static_assert(QuantGroupSize::kN == 1, "no block for N supported yet!"); + // static_assert(QuantGroupSize::kN == 1, "no block for N supported yet!"); using I0 = number<0>; using I1 = number<1>; diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp index 067415526c..7f1f4a4d48 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp @@ -171,10 +171,10 @@ template struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding_pattern { - // TODO: make pattern where below condition does not need to hold - GGemmMultiDSplitk! static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!"); static constexpr index_t warp_size = get_warp_size(); static constexpr index_t num_warps = BlockSize / get_warp_size(); @@ -186,36 +186,49 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding static constexpr index_t NIterPerWarp = BlockGemmShape::kN / (NWarps * WarpGemm::kN); static_assert(num_warps == MWarps * NWarps * KWarps); - - // KWarps > 1 isn't supported static_assert(KWarps == 1); - // # of elements per thread - static constexpr index_t X = XPerTile; - static constexpr index_t XR = 2; - - // Number of iters per warp - // MIters are indexed using (Y0, Y1) - // TODO: does this need to change? X/Y-PerTile are already the smaller size. Is NIterPerWarp good for group size N>1? - // TODO: should we split it so that one of the dims is the number of scales within the block? -> we should have less data per thread - static constexpr index_t Y0 = NIterPerWarp; - - // # of warps in Y dim - static constexpr index_t Y1 = NWarps; - - static constexpr index_t Y2 = WarpGemm::kN; - - static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover the blocktile along Y."); - CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution() { - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<0, 1>>, - tuple, sequence<1, 2>>, - sequence<1, 2>, - sequence<0, 0>>{}); + if constexpr(YPerQ == 1) + { + // YPerQ == 1 implementation + constexpr index_t X = XPerTile; + constexpr index_t XR = 2; + constexpr index_t Y0 = NIterPerWarp; + constexpr index_t Y1 = NWarps; + constexpr index_t Y2 = WarpGemm::kN; + + static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover the blocktile along Y."); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<0, 1>>, + tuple, sequence<1, 2>>, + sequence<1, 2>, + sequence<0, 0>>{}); + } + else + { + // YPerQ > 1 implementation + // TODO: do not repeat everything to all threads + + // return make_static_tile_distribution( + // tile_distribution_encoding, + // tuple, sequence>, + // tuple, sequence<0, 1>>, + // tuple, sequence<1, 1>>, + // sequence<1, 2>, + // sequence<0, 0>>{}); + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<0>>, + tuple, sequence<0>>, + sequence<1, 2>, + sequence<0, 0>>{}); + } } }; 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 a1fefda056..dfe99f6648 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 @@ -383,7 +383,8 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBaseis_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{}))); + ck_tile::HostTensor bq_bqk_bqn( + ck_tile::host_tensor_descriptor(BQK, BQN, 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); - + ck_tile::FillUniformDistribution{0.001f, 0.01f}(bq_bqk_bqn); // 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 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)); // Copy to device @@ -426,7 +426,7 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase(a_m_k, bq_bqk_n, b_k_n, c_m_n_host_ref); + false>(a_m_k, bq_bqk_bqn, b_k_n, c_m_n_host_ref); // Get device result ck_tile::HostTensor c_m_n_dev_result( 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 index 46d37d0dfc..89306d8dc6 100644 --- a/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp @@ -21,6 +21,8 @@ using BQuantGrouped = std::integral_constant; using TensorQuant = std::integral_constant; using GroupSize = ck_tile::QuantGroupShape>; +using GroupSize64 = ck_tile::QuantGroupShape>; +using GroupSize2D = ck_tile::QuantGroupShape>; // Type combinations for each quantization type // clang-format off @@ -31,47 +33,56 @@ using AQuantTypes = ::testing::Types< std::tuple, std::tuple, - // PreshuffleQuant = false && TransposeC = true - std::tuple, - std::tuple, - std::tuple, - std::tuple, + // // PreshuffleQuant = false && TransposeC = true + // std::tuple, + // std::tuple, + // std::tuple, + // std::tuple, - // PreshuffleQuant = true && TransposeC = false - std::tuple, - std::tuple, - std::tuple, - std::tuple, + // // PreshuffleQuant = true && TransposeC = false + // std::tuple, + // std::tuple, + // std::tuple, + // std::tuple, - // PreshuffleQuant = true && TransposeC = true - std::tuple, - std::tuple, - std::tuple, + // // PreshuffleQuant = true && TransposeC = true + // std::tuple, + // std::tuple, + // std::tuple, std::tuple >; // clang-format on // clang-format off using BQuantTypes = ::testing::Types< + // 1d cases with grouping only on k axis std::tuple, - std::tuple, - std::tuple, - std::tuple + // std::tuple, + // std::tuple, + // std::tuple, + + // std::tuple, + // std::tuple, + // std::tuple, + // std::tuple + // 2d cases with grouping also on the n axis + std::tuple >; // clang-format on // clang-format off using BPreshuffleBQuantTypes = ::testing::Types< std::tuple, - std::tuple, - std::tuple, - std::tuple, + // std::tuple, + // std::tuple, + // std::tuple, - std::tuple, - std::tuple, - std::tuple, + // std::tuple, + // std::tuple, + // std::tuple, std::tuple >; +// clang-format on // clang-format off using RowColQuantTypes = ::testing::Types<