ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 09:38:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add more specialized tile distributions

Browse Source
This commit is contained in:
Sami Remes
2025-10-27 13:43:02 +00:00
parent d100ab690a
commit 37738e4cb8
3 changed files with 44 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp
View File

@@ -343,21 +343,23 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
// Multiply bquant with accumulated C
const index_t reg_offset = [&]() {
if constexpr(Traits::NQPerBlock == Traits::NPerBlock)
constexpr bool scale_per_niter_per_warp = Traits::QuantGroupSize::kN == 1 || Traits::NQPerBlock >= Traits::NIterPerWarp * Traits::NWarp;
if constexpr(scale_per_niter_per_warp)
{
// Each row of B has a separate scale, each thread has its own
// single element of the scale matrix for the current nIter/kQScale
// Each nIter and warp/thread has its own scale - tile dstr handles the proper loading
return nIter * Traits::BQPerBlock + kQScale;
}
else
{
// FIXME: temporarily the tile distribution replicates all block's
// scales to all threads -> need to calculate the index manually
// here from nIter and warp id
// Many warps/iters can share the same scale, index from full [NQPerBlock, BQPerBlock] matrix
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(get_lane_id() == 0)
{
printf("row_index: %d\n", row_index);
}
return row_index * Traits::BQPerBlock + kQScale;
}
}();

4
include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp
View File

@@ -29,8 +29,8 @@ struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase<Prob
static constexpr index_t NPerBlockBQ = NPerBlock / QuantGroupSize::kN;
static constexpr index_t KPerBlockBQ = KPerBlock / QuantGroupSize::kK;
static_assert(NPerBlockBQ > 1, "NPerBlock must be > QuantGroupSize");
static_assert(KPerBlockBQ > 1, "KPerBlock must be > QuantGroupSize");
static_assert(NPerBlockBQ >= 1, "NPerBlock must be >= QuantGroupSize");
static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= QuantGroupSize");
static_assert(NPerBlock % QuantGroupSize::kN == 0,
"NPerBlock must be a multiple of QuantGroupSize");

36
include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp
View File

@@ -209,10 +209,37 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
sequence<1, 2>,
sequence<0, 0>>{});
}
else if constexpr(YPerTile >= NIterPerWarp * NWarps)
{
// small NQ block size case: split NQ axis by iters and Nwarps
constexpr index_t NQPerIter = integer_divide_ceil(YPerTile, NIterPerWarp * NWarps);
constexpr index_t XR = get_warp_size() / NQPerIter;
static_assert(YPerTile == NQPerIter * NWarps * NIterPerWarp);
return make_static_tile_distribution(
tile_distribution_encoding<sequence<MWarps, XR>,
tuple<sequence<NIterPerWarp, NWarps, NQPerIter>, sequence<XPerTile>>,
tuple<sequence<0, 1>, sequence<0, 1>>,
tuple<sequence<0, 1>, sequence<1, 2>>,
sequence<1, 2>,
sequence<0, 0>>{});
}
// else if constexpr(YPerTile >= NIterPerWarp)
// {
// // now all NWarps have the same scale -> replicate
// constexpr index_t NQPerIter = integer_divide_ceil(YPerTile, NIterPerWarp);
// constexpr index_t XR = get_warp_size() / NQPerIter;
// static_assert(YPerTile == NQPerIter * NWarps * NIterPerWarp);
// return make_static_tile_distribution(
// tile_distribution_encoding<sequence<MWarps, NWarps, XR>,
// tuple<sequence<NIterPerWarp, NQPerIter>, sequence<XPerTile>>,
// tuple<sequence<0, 0>, sequence<0, 1>>,
// tuple<sequence<0, 1>, sequence<2, 1>>,
// sequence<1, 2>,
// sequence<0, 0>>{});
// }
else
{
// YPerQ > 1 implementation - each group of YPerQ rows share the same scale
// TODO: do not repeat everything to all threads
// larger NQ block size, multiple iters/warps use same scales -> replicate to all threads
return make_static_tile_distribution(
tile_distribution_encoding<sequence<MWarps, NWarps, get_warp_size()>,
tuple<sequence<YPerTile>, sequence<XPerTile>>,
@@ -222,6 +249,11 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
sequence<0, 0>>{});
}
}
CK_TILE_HOST_DEVICE static constexpr bool is_fully_replicated()
{
return YPerQ > 1 && YPerTile < NIterPerWarp * NWarps;
}
};
template <typename GroupSizes>