diff --git a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_base.hpp index 316e52b860..5d6f7ec16a 100644 --- a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_base.hpp +++ b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_base.hpp @@ -48,7 +48,7 @@ struct GemmMXPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase); using YPerTile = number; - using XPerTile = number; + using XPerTile = number; auto a_copy_draw_window = make_tile_window(a_scale_dram_block_window_tmp.get_bottom_tensor_view(), @@ -66,7 +66,7 @@ struct GemmMXPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase); using YPerTile = number; - using XPerTile = number; + using XPerTile = number; auto b_copy_draw_window = make_tile_window(b_scale_dram_block_window_tmp.get_bottom_tensor_view(), diff --git a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_policy.hpp index 36cb705481..020f15ad3c 100644 --- a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_policy.hpp @@ -18,6 +18,10 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol using Base::ATileAccessPattern; using Base::BTileAccessPattern; + static constexpr auto MXdlPack = 2; + static constexpr auto NXdlPack = 2; + static constexpr auto KXdlPack = 2; + template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeAScale() { @@ -44,37 +48,6 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol return GetScaleGlobalVectorLoadSize(); } - template - CK_TILE_HOST_DEVICE static constexpr auto MakeAQDramTileDistribution() - { - using AQLayout = remove_cvref_t; - using BlockGemmShape = typename Problem::BlockGemmShape; - - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPerBlockAQ = KPerBlock / Problem::kQuantGroupSize; - constexpr index_t VecLoadSize = GetVectorSizeAQ(); - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmMfmaDispatcher; - - static_assert(std::is_same_v); - using TileEncodingPattern = TileDistributionEncodingPatternAQ; - - return TileEncodingPattern::Make2DStaticTileDistribution(); - } - // A Scale DRAM tile distribution // This is used to load the A scale data from DRAM into shared memory. template @@ -82,32 +55,30 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol { using AScaleLayout = remove_cvref_t; using BlockGemmShape = typename Problem::BlockGemmShape; + using AScaleDataType = remove_cvref_t; + static constexpr index_t APackedSize = + ck_tile::numeric_traits>::PackedSize; constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t MPerBlockScale = Problem::BlockGemmShape::kM / MXdlPack; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPerBlockScale = KPerBlock / Problem::kBlockScaleSize; - constexpr index_t VecLoadSize = GetVectorSizeAScale(); - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmMfmaDispatcher; - static constexpr auto MXdlPack = 2; - static constexpr auto NXdlPack = 2; - static constexpr auto KXdlPack = 2; static_assert(std::is_same_v); using TileEncodingPattern = TileDistributionEncodingPatternAScale; + MPerBlockScale, + KPerBlockScale>; return TileEncodingPattern::Make2DStaticTileDistribution(); } @@ -116,6 +87,33 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol template CK_TILE_HOST_DEVICE static constexpr auto MakeBScaleDramTileDistribution() { + using BScaleLayout = remove_cvref_t; + using BlockGemmShape = typename Problem::BlockGemmShape; + using BScaleDataType = remove_cvref_t; + static constexpr index_t BPackedSize = + ck_tile::numeric_traits>::PackedSize; + + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t NPerBlockScale = Problem::BlockGemmShape::kN / NXdlPack; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t KPerBlockScale = + KPerBlock * BPackedSize / (roblem::kBlockScaleSize * KXdlPack); + using WarpTile = typename Problem::BlockGemmShape::WarpTile; + using WarpGemm = WarpGemmMfmaDispatcher; + + static_assert(std::is_same_v); + using TileEncodingPattern = TileDistributionEncodingPatternBScale; + return TileEncodingPattern::Make2DStaticTileDistribution(); } template diff --git a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_utils.hpp b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_utils.hpp index 1fda95e6aa..db88b6fb0d 100644 --- a/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_utils.hpp +++ b/include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_utils.hpp @@ -42,56 +42,6 @@ CK_TILE_HOST_DEVICE static constexpr auto GetScaleGlobalVectorLoadSize() return PackedSize; // Absolute fallback } -// AQ holds groupquant scale data for A. Data is loaded from DRAM and partitioned across -// threads. Post mfma scales are shuffled across threads in the warp and applied to -// accum registers. -template -struct TileDistributionEncodingPatternAQ : public TileDistributionEncodingPattern -{ - // 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(); - - static constexpr index_t MWarps = BlockGemmShape::BlockWarps::at(number<0>{}); - static constexpr index_t NWarps = BlockGemmShape::BlockWarps::at(number<1>{}); - static constexpr index_t KWarps = BlockGemmShape::BlockWarps::at(number<2>{}); - - static constexpr index_t MIterPerWarp = BlockGemmShape::kM / (MWarps * WarpGemm::kM); - - 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 Y0 = 1; - static constexpr index_t Y1 = MIterPerWarp ? MIterPerWarp : 1; - static constexpr index_t Y2 = MWarps; - static constexpr index_t Y3 = WarpGemm::kM; - static_assert(Y3 >= WarpGemm::kM, "Scales for all rows must be available within the warp."); - static_assert(Y0 * Y1 * Y2 * Y3 == YPerTile, - "Y0, Y1, Y2, Y3 must cover the blocktile along Y."); - - CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution() - { - return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 1>>, - tuple, sequence<0, 3>>, - sequence<1, 2>, - sequence<1, 0>>{}); - } -}; - // A Scale data for A data is preshuffled and loaded from DRAM // using v_mfama_f32_scale_f32_16x16x128_F8F6F4 instruction for calculating template struct TileDistributionEncodingPatternAScale : public TileDistributionEncodingPattern { @@ -120,16 +67,15 @@ struct TileDistributionEncodingPatternAScale : public TileDistributionEncodingPa static_assert(KWarps == 1, "KWarps > 1 is not supported"); // Y dimension (M) decomposition - static constexpr index_t MXdlPack = 2; // MXdlPack is always 2 - static constexpr index_t Y1 = MWarps; - static constexpr index_t Y2 = MThreadPerXdl; - static constexpr index_t Y0 = YPerTile / MXdlPack / (MWarps * MThreadPerXdl); + static constexpr index_t Y1 = MWarps; + static constexpr index_t Y2 = MThreadPerXdl; + static constexpr index_t Y0 = YPerTile / (MWarps * MThreadPerXdl); // X dimension (K) decomposition static constexpr index_t X0 = KThreadPerXdl; static constexpr index_t X1 = VecSize; - static_assert(Y0 * Y1 * Y2 * Y3 == YPerTile, "Y dimensions must cover the YPerTile"); + static_assert(Y0 * Y1 * Y2 == YPerTile, "Y dimensions must cover the YPerTile"); static_assert(X0 * X1 == XPerTile, "X dimensions must cover the XPerTile"); CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution() @@ -137,8 +83,57 @@ struct TileDistributionEncodingPatternAScale : public TileDistributionEncodingPa return make_static_tile_distribution( tile_distribution_encoding, tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 2>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } +}; + +// B Scale data for B data is preshuffled and loaded from DRAM +// using v_mfama_f32_scale_f32_16x16x128_F8F6F4 instruction for calculating +template +struct TileDistributionEncodingPatternBScale : public TileDistributionEncodingPattern +{ + static_assert(NPerBlock % NXdlPack == 0, "XPerTile must be a multiple of VecSize!"); + static constexpr index_t warp_size = get_warp_size(); + static constexpr index_t num_warps = BlockSize / warp_size; + + static constexpr index_t MWarps = BlockGemmShape::BlockWarps::at(number<0>{}); + static constexpr index_t NWarps = BlockGemmShape::BlockWarps::at(number<1>{}); + static constexpr index_t KWarps = BlockGemmShape::BlockWarps::at(number<2>{}); + + static constexpr index_t NThreadPerXdl = WarpGemm::kN; + static constexpr index_t KThreadPerXdl = warp_size / NThreadPerXdl; + + static_assert(num_warps == MWarps * NWarps * KWarps, "Block warps do not match block size"); + static_assert(KWarps == 1, "KWarps > 1 is not supported"); + + // Y dimension (N) decomposition + static constexpr index_t Y1 = NWarps; + static constexpr index_t Y2 = NThreadPerXdl; + static constexpr index_t Y0 = YPerTile / (NWarps * NThreadPerXdl); + + // X dimension (K) decomposition + static constexpr index_t X0 = KThreadPerXdl; + static constexpr index_t X1 = VecSize; + + static_assert(Y0 * Y1 * Y2 == YPerTile, "Y dimensions must cover the YPerTile"); + static_assert(X0 * X1 == XPerTile, "X dimensions must cover the XPerTile"); + + CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution() + { + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<0, 2>>, sequence<1, 2>, sequence<0, 1>>{}); }