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

formatting

Browse Source
This commit is contained in:
khushbu
2025-12-11 22:29:23 -05:00
parent 995d1a5cf6
commit 44aaaacbec
8 changed files with 184 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -349,6 +349,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
constexpr index_t reg_offset = nIter;
auto pull_from_lane =
(__lane_id() & (WarpGemm::kN - 1)) * Traits::KQPerBlock + kQScale;
auto& scale_reg = bq_block_tensor.get_thread_buffer()[reg_offset];
// cross lane ops
uint32_t scale_reg_dword;
@@ -368,19 +369,20 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
float scale_reg_f = Base::cvt_scale_to_fp32(gathered_scale_reg);
printf("block_id: %d, warp_id: %d, thread_id(): %d, nIter: %d, lane_id(): "
"%u, kQScale: %d, pull_from_lane: %u, scale_reg: %f, "
"scale_reg_f: %f\n",
get_block_id(),
get_warp_id(),
get_thread_id(),
static_cast<int>(nIter),
__lane_id(),
static_cast<int>(kQScale),
pull_from_lane,
scale_reg,
scale_reg_f);
// printf("block_id: %d, warp_id: %d, thread_id(): %d, nIter: %d,
// lane_id(): "
// "%u, kQScale: %d, pull_from_lane: %u, scale_reg: %f, "
// "scale_reg_f: %f\n",
// get_block_id(),
// get_warp_id(),
// get_thread_id(),
// static_cast<int>(nIter),
// __lane_id(),
// static_cast<int>(kQScale),
// pull_from_lane,
// scale_reg,
// scale_reg_f);
static_for<0, WarpGemm::kM * WarpGemm::kN / warp_size, 1>{}(
[&](auto c_row) {
c_block_tensor.get_thread_buffer()[tbuf_offset + c_row] +=

116
include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp
View File

@@ -298,15 +298,15 @@ struct QuantGemmKernel
const auto bq_x = N * KPerBlockBQ; // 2x2 = 4
const auto bq_y = QK_B / KPerBlockBQ; // 4/2 = 2
if(get_block_id() == 0 && get_thread_id() == 0)
{
printf("N:%d, QK_B:%d\n", N, QK_B);
printf("bq_x: %d, bq_y: %d, getVectorSizeBQ: %d, kPerBlockBQ: %d\n",
bq_x,
bq_y,
GetVectorSizeBQ,
KPerBlockBQ);
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("N:%d, QK_B:%d\n", N, QK_B);
// printf("bq_x: %d, bq_y: %d, getVectorSizeBQ: %d, kPerBlockBQ: %d\n",
// bq_x,
// bq_y,
// GetVectorSizeBQ,
// KPerBlockBQ);
// }
const auto bq_desc = make_naive_tensor_descriptor(make_tuple(bq_y, bq_x),
make_tuple(bq_x, 1),
@@ -319,10 +319,10 @@ struct QuantGemmKernel
// each thread block can process complete tiles without edge cases
const auto block_tile_size = NPerBlockBQ * KPerBlockBQ; // 2x2 = 4
if(get_block_id() == 0 && get_thread_id() == 0)
{
printf("block_tile_size:%d \n", block_tile_size);
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("block_tile_size:%d \n", block_tile_size);
// }
const auto bq_pad0_desc = transform_tensor_descriptor(
bq_desc,
@@ -337,22 +337,25 @@ struct QuantGemmKernel
// Split the X dimension into [wave_tile_count_x, wave_tile_size]
// This separates the work into tiles that can be processed by
// individual warps/waves
const auto pad_bq_x = bq_pad0_desc.get_lengths()[I1]; // 4
const auto wave_tile_size = ((QN_B <= WarpTileN) ? (WarpTileN / QN_B) : 1) * KPerBlockBQ; // 32/16 x 2 = 4 = 2
const auto pad_bq_x = bq_pad0_desc.get_lengths()[I1]; // 4
const auto wave_tile_size =
((QN_B <= WarpTileN) ? (WarpTileN / QN_B) : 1 /*QN_B/WarpTileN*/) *
KPerBlockBQ; // 32/16 x 2 = 4 = 2
const auto wave_tile_count_x =
ck_tile::integer_divide_ceil(pad_bq_x, wave_tile_size); // 4/4 = 1 ==2
if(get_block_id() == 0 && get_thread_id() == 0)
{
printf("pad_bq_x:%d, WarpTileN:%d, NPerBlockPQ: %d, KPerBlockBQ: %d, wave_tile_size: "
"%d, wave_tile_count_x: %d\n",
pad_bq_x,
WarpTileN,
NPerBlockBQ,
KPerBlockBQ,
wave_tile_size,
wave_tile_count_x);
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("pad_bq_x:%d, WarpTileN:%d, NPerBlockPQ: %d, KPerBlockBQ: %d, wave_tile_size:
// "
// "%d, wave_tile_count_x: %d\n",
// pad_bq_x,
// WarpTileN,
// NPerBlockBQ,
// KPerBlockBQ,
// wave_tile_size,
// wave_tile_count_x);
// }
const auto bq_unmerge_pad0_desc = transform_tensor_descriptor(
bq_pad0_desc,
@@ -383,13 +386,16 @@ struct QuantGemmKernel
// where merged_outer_dim = bq_y * wave_tile_count_x
// This layout facilitates efficient block-to-data mapping
const auto pad_wave_size = ck_tile::integer_least_multiple(wave_tile_size, get_warp_size());
if(get_block_id() == 0 && get_thread_id() == 0)
{
printf("pad_wave_size:%d\n", pad_wave_size);
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("pad_wave_size:%d\n", pad_wave_size);
// printf("Final bq tensor lengths: %d x %d \n",
// bq_y * wave_tile_count_x,
// pad_wave_size);
// }
const auto bq_merge_pad1_desc = transform_tensor_descriptor(
bq_pad1_desc,
make_tuple(make_merge_transform(make_tuple(bq_y, wave_tile_count_x)), // 2
make_tuple(make_merge_transform(make_tuple(bq_y, wave_tile_count_x)), // 4
make_pass_through_transform(pad_wave_size)), // 64
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -1115,13 +1121,33 @@ struct QuantGemmKernel
if constexpr(PreshuffleQuant)
{
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
constexpr auto block_n = TilePartitioner::NPerBlock / QuantGroupSize::kN;
constexpr auto warp_n = TilePartitioner::BlockGemmShape::WarpTile::at(I1);
constexpr auto block_n =
TilePartitioner::NPerBlock / QuantGroupSize::kN; // 64 / 32 = 2
constexpr auto warp_n = TilePartitioner::BlockGemmShape::WarpTile::at(I1);
constexpr auto bqk_per_block = TilePartitioner::KPerBlock / QuantGroupSize::kK;
constexpr auto tile_window_width = ck_tile::integer_least_multiple(
warp_n * bqk_per_block, get_warp_size()); // 128
constexpr auto tile_window_height = block_n / warp_n; // 2
auto block_n_idx = i_n / block_n;
warp_n * bqk_per_block, get_warp_size()); // 128
constexpr auto tile_window_height =
min(block_n,
TilePartitioner::BlockGemmShape::BlockWarps::at(
I1)); // block_n / warp_n; // 2 / 4 = 0
auto block_n_idx = i_n / TilePartitioner::NPerBlock; // 0,1,2
// if(get_thread_id() == 0)
// {
// printf("In MakeGemmTileWindows for BQ with PreshuffleQuant\n");
// printf("block_id: %d, block_n: %d, warp_n: %d, bqk_per_block: %d,
// block_n_idx: %d, "
// "tile_window_width: %d, tile_window_height: %d, i_n: %d\n",
// get_block_id(),
// static_cast<int>(block_n),
// static_cast<int>(warp_n),
// static_cast<int>(bqk_per_block),
// static_cast<int>(block_n_idx),
// tile_window_width,
// static_cast<int>(tile_window_height),
// static_cast<int>(i_n));
// }
return make_tile_window(
bq_pad_view,
@@ -1226,15 +1252,15 @@ struct QuantGemmKernel
{
n = kargs.N;
}
if(get_block_id() == 0 && get_thread_id() == 0)
{
printf("In RunGemm, before GemmPipeline call for BQuantGrouped\n");
// To print Tile window after bq_pad0_desc
// bq_block_window.template print_tile_window_range<BQDataType>(
// 0, 128, 0, 2, "bq block window");
bq_block_window.template print_tile_window_range<BQDataType>(
0, 8, 0, 64, "bq block window");
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// printf("In RunGemm, before GemmPipeline call for BQuantGrouped\n");
// // To print Tile window after bq_pad0_desc
// // bq_block_window.template print_tile_window_range<BQDataType>(
// // 0, 128, 0, 2, "bq block window");
// bq_block_window.template print_tile_window_range<BQDataType>(
// 0, 8, 0, 64, "bq block window");
// }
return GemmPipeline{}.template operator()(
a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr_0, n);
}

11
include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp
View File

@@ -196,7 +196,7 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3<Prob
const BDramBlockWindowTmp& b_dram_block_window_tmp,
const BElementFunction& b_element_func,
const BQDramBlockWindowTmp& bq_dram_block_window_tmp,
index_t n,
[[maybe_unused]] index_t n,
index_t num_loop,
void* p_smem) const
{
@@ -280,9 +280,12 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3<Prob
constexpr BDramTileWindowStep b_dram_tile_window_step =
is_b_row_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock);
const BQDramTileWindowStep bq_dram_tile_window_step =
(PreshuffleQuant) ? make_array(ck_tile::integer_least_multiple(n, NPerBlock) /
BlockGemmShape::WarpTile::at(number<1>{}),
0)
(PreshuffleQuant)
? make_array(((NPerBlockBQ < BlockGemmShape::BlockWarps::at(number<1>{}))
? ck_tile::integer_divide_ceil(n, QuantGroupSize::kN)
: ck_tile::integer_least_multiple(n, NPerBlock) /
BlockGemmShape::WarpTile::at(number<1>{})),
0)
: is_bq_row_major ? make_array(KPerBlockBQ, 0)
: make_array(0, KPerBlockBQ);

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

@@ -245,56 +245,70 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
if constexpr(NPerQ <= WarpGemm::kN)
{
constexpr auto N1 = BlockGemmShape::kK / KPerQ; // 2
constexpr auto N0 = WarpGemm::kN / NPerQ; //BlockGemmShape::kN / KPerQ; // 1
constexpr auto N0 = WarpGemm::kN / NPerQ; // BlockGemmShape::kN / KPerQ; // 1
constexpr auto N2 = 1;
constexpr auto NR1 = NPerQ; // 16
constexpr auto NR1 = NPerQ; // 16
constexpr auto NR0 = warp_size / (N0 * N1 * N2 * NR1); // 64/(1*2*1*16)=2
constexpr auto K1 = NWarps; // 4
constexpr auto K0 = KPerTile / K1; // 1
constexpr auto K1 = NWarps; // 4
constexpr auto K0 = KPerTile / K1; // 1
constexpr auto KR = 1;
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<MWarps, NR0, NR1, KR>,
tuple<sequence<K0, K1>, sequence<N0, N1, N2>>,
tuple<sequence<0, 1>, sequence<0, 2, 0, 2, 0>>, // (Mwarp, Nwarp),(XR0, X0, XR1, X1, YR)
tuple<sequence<0, 1>, sequence<0, 2, 0, 2, 0>>, // (Mwarp, Nwarp),(XR0, X0,
// XR1, X1, YR)
tuple<sequence<0, 1>, sequence<1, 0, 2, 1, 3>>, // (1, 4), (2, 1, 16, 2, 1)
sequence<1, 2>,
sequence<0, 2>>{});
}
else if constexpr(NPerQ <= WarpGemm::kN * NWarps)
else if constexpr(NPerQ < WarpGemm::kN * NWarps)
{
constexpr auto KR = NPerQ / WarpGemm::kN; // Scale replication factor 32/16 = 2
constexpr auto K1 = NWarps / KR; // Warps per unique scale 4/2 = 2
constexpr auto K0 = KPerTile / K1; // Iterations to cover N dimension 4/2 = 2
constexpr auto KR = NPerQ / WarpGemm::kN; // Scale replication factor 64/16 = 4
constexpr auto K1 = NWarps / KR; // Warps per unique scale 4/4 = 1
constexpr auto K0 = KPerTile / K1; // Iterations to cover N dimension 4/1 = 4
constexpr auto N1 = BlockGemmShape::kK / KPerQ; // 2
constexpr auto N0 = 1; //NPerQ/WarpGemm::kN; // 2
constexpr auto N0 = 1; // NPerQ/WarpGemm::kN; // 1
constexpr auto N2 = 1;
constexpr auto NR1 = NPerQ; // 32
constexpr auto NR1 = NPerQ; // 32
constexpr auto NR0 = warp_size / (N0 * N1 * N2 * NR1); // 64/(1*2*1*32)=1
if(get_block_id() == 0 && get_thread_id() == 0)
{
// Debug print to verify values
printf("PreshuffleQuant Medium-grained: MWarps: %d, K1=%d, KR=%d, get_warp_size(): %d, K0=%d, N0=%d\n",
MWarps,
K1,
KR,
get_warp_size(),
K0,
N0);
}
// if(get_block_id() == 0 && get_thread_id() == 0)
// {
// // Debug print to verify values
// printf("PreshuffleQuant Medium-grained: MWarps: %d, K1=%d, KR=%d,
// get_warp_size(): %d, K0=%d, N0=%d\n",
// MWarps,
// K1,
// KR,
// get_warp_size(),
// K0,
// N0);
// }
return make_static_tile_distribution(
tile_distribution_encoding<sequence<MWarps, NR0, NR1, KR>,
tuple<sequence<K0, K1>, sequence<N0, N1, N2>>,
tuple<sequence<0, 1, 0>, sequence<0, 2, 0, 2>>,
tuple<sequence<0, 1, 3>, sequence<1, 0, 2, 1>>,
sequence<1, 2>,
sequence<0, 2>>{});
tuple<sequence<K0, K1>, sequence<N0, N1, N2>>,
tuple<sequence<0, 1, 0>, sequence<0, 2, 0, 2>>,
tuple<sequence<0, 1, 3>, sequence<1, 0, 2, 1>>,
sequence<1, 2>,
sequence<0, 2>>{});
}
else
{
constexpr auto N1 = BlockGemmShape::kK / KPerQ; // 2
constexpr auto N0 = 1; // NPerQ/WarpGemm::kN; // 1
constexpr auto N2 = 1;
constexpr auto NR1 = 32; // 32
constexpr auto NR0 = warp_size / (N0 * N1 * N2 * NR1); // 64/(1*2*1*32)=1
return make_static_tile_distribution(
tile_distribution_encoding<sequence<MWarps, NWarps, NR0, NR1>,
tuple<sequence<KPerTile>, sequence<N0, N1, N2>>,
tuple<sequence<0, 0>, sequence<0, 2, 0, 2>>,
tuple<sequence<0, 1>, sequence<2, 0, 3, 1>>,
sequence<1, 2>,
sequence<0, 2>>{});
}
}
else