ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

updated A, B scale distribution encoding patterns

Browse Source
This commit is contained in:
mtgu0705
2025-08-13 04:42:14 -05:00
parent 5e5b9bdbd0
commit 1698930818
3 changed files with 98 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_base.hpp
View File

@@ -48,7 +48,7 @@ struct GemmMXPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase<Problem,
static_assert(
std::is_same_v<typename Problem::AScaleLayout, tensor_layout::gemm::RowMajor>);
using YPerTile = number<MPerBlock / MXdlPack>;
using XPerTile = number<KPerBlock / BlockScaleSize / KXdlPack>;
using XPerTile = number<KPerBlock * APackedSize / (BlockScaleSize * KXdlPack)>;
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<Problem,
static_assert(
std::is_same_v<typename Problem::BScaleLayout, tensor_layout::gemm::ColumnMajor>);
using YPerTile = number<NPerBlock / NXdlPack>;
using XPerTile = number<KPerBlock / BlockScaleSize / KXdlPack>;
using XPerTile = number<KPerBlock * BPackedSize / (BlockScaleSize * KXdlPack)>;
auto b_copy_draw_window =
make_tile_window(b_scale_dram_block_window_tmp.get_bottom_tensor_view(),

84
include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_pipeline_ag_bg_cr_policy.hpp
View File

@@ -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 <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeAScale()
{
@@ -44,37 +48,6 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol
return GetScaleGlobalVectorLoadSize<Problem, BScaleDataType, NPerBlock, KPerBlockScale>();
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeAQDramTileDistribution()
{
using AQLayout = remove_cvref_t<typename Problem::AQLayout>;
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<Problem>();
using WarpTile = typename Problem::BlockGemmShape::WarpTile;
using WarpGemm = WarpGemmMfmaDispatcher<typename Problem::ComputeDataType,
typename Problem::ComputeDataType,
typename Problem::CDataType,
WarpTile::at(I0),
WarpTile::at(I1),
WarpTile::at(I2),
false>;
static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
using TileEncodingPattern = TileDistributionEncodingPatternAQ<BlockGemmShape,
WarpGemm,
BlockSize,
MPerBlock,
KPerBlockAQ,
VecLoadSize>;
return TileEncodingPattern::Make2DStaticTileDistribution();
}
// A Scale DRAM tile distribution
// This is used to load the A scale data from DRAM into shared memory.
template <typename Problem>
@@ -82,32 +55,30 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol
{
using AScaleLayout = remove_cvref_t<typename Problem::AScaleLayout>;
using BlockGemmShape = typename Problem::BlockGemmShape;
using AScaleDataType = remove_cvref_t<typename Problem::AScaleDataType>;
static constexpr index_t APackedSize =
ck_tile::numeric_traits<remove_cvref_t<ADataType>>::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<Problem>();
using WarpTile = typename Problem::BlockGemmShape::WarpTile;
using WarpGemm = WarpGemmMfmaDispatcher<typename Problem::ComputeDataType,
constexpr index_t KPerBlockScale =
KPerBlock * APackedSize / (Problem::kBlockScaleSize * KXdlPack);
using WarpTile = typename Problem::BlockGemmShape::WarpTile;
using WarpGemm = WarpGemmMfmaDispatcher<typename Problem::ComputeDataType,
typename Problem::ComputeDataType,
typename Problem::CDataType,
WarpTile::at(I0),
WarpTile::at(I1),
WarpTile::at(I2),
false>;
static constexpr auto MXdlPack = 2;
static constexpr auto NXdlPack = 2;
static constexpr auto KXdlPack = 2;
static_assert(std::is_same_v<AScaleLayout, tensor_layout::gemm::RowMajor>);
using TileEncodingPattern = TileDistributionEncodingPatternAScale<BlockGemmShape,
WarpGemm,
BlockSize,
MPerBlock,
KPerBlockScale,
MXdlPack,
KXdlPack>;
MPerBlockScale,
KPerBlockScale>;
return TileEncodingPattern::Make2DStaticTileDistribution();
}
@@ -116,6 +87,33 @@ struct GemmMXPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgCrPol
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeBScaleDramTileDistribution()
{
using BScaleLayout = remove_cvref_t<typename Problem::BScaleLayout>;
using BlockGemmShape = typename Problem::BlockGemmShape;
using BScaleDataType = remove_cvref_t<typename Problem::BScaleDataType>;
static constexpr index_t BPackedSize =
ck_tile::numeric_traits<remove_cvref_t<BDataType>>::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<typename Problem::ComputeDataType,
typename Problem::ComputeDataType,
typename Problem::CDataType,
WarpTile::at(I0),
WarpTile::at(I1),
WarpTile::at(I2),
false>;
static_assert(std::is_same_v<BScaleLayout, tensor_layout::gemm::ColumnMajor>);
using TileEncodingPattern = TileDistributionEncodingPatternBScale<BlockGemmShape,
WarpGemm,
BlockSize,
NPerBlockScale,
KPerBlockScale>;
return TileEncodingPattern::Make2DStaticTileDistribution();
}
template <typename Problem>

115
include/ck_tile/ops/gemm_mx/pipeline/gemm_mx_utils.hpp
View File

@@ -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 <typename BlockGemmShape,
typename WarpGemm,
index_t BlockSize,
index_t YPerTile,
index_t XPerTile,
index_t VecSize>
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<sequence<NWarps>,
tuple<sequence<Y0, Y1, Y2, Y3>, sequence<X>>,
tuple<sequence<1, 0>, sequence<1, 1>>,
tuple<sequence<2, 0>, 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 <typename BlockGemmShape,
@@ -99,9 +49,6 @@ template <typename BlockGemmShape,
index_t BlockSize,
index_t YPerTile,
index_t XPerTile,
index_t MXdlPack,
index_t NXdlPack,
index_t KXdlPack,
index_t VecSize = 1>
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<sequence<NWarps>,
tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
tuple<sequence<1, 0>, sequence<2, 1>>,
tuple<sequence<1, 0>, 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 <typename BlockGemmShape,
typename WarpGemm,
index_t BlockSize,
index_t YPerTile,
index_t XPerTile,
index_t NXdlPack,
index_t VecSize = 1>
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<sequence<MWarps>,
tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
tuple<sequence<0, 1>, sequence<2, 1>>,
tuple<sequence<0, 1>, sequence<0, 2>>,
sequence<1, 2>,
sequence<0, 1>>{});
}