ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 02:05:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP

Browse Source
This commit is contained in:
Sami Remes
2025-10-20 15:42:39 +00:00
parent 9988a46af2
commit 36b88c665c
4 changed files with 28 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -345,19 +345,20 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase<Problem_>
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
// 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)
// 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,

15
include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp
View File

@@ -51,13 +51,14 @@ struct GemmBQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC
Problem::TransposeC>;
static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
using TileEncodingPattern = tile_distribution_encoding_pattern_bq<BlockGemmShape,
WarpGemm,
BlockSize,
NPerBlockBQ,
KPerBlockBQ,
Problem::QuantGroupSize::kN,
VecLoadSize>;
using TileEncodingPattern =
tile_distribution_encoding_pattern_bq<BlockGemmShape,
WarpGemm,
BlockSize,
NPerBlockBQ,
KPerBlockBQ,
Problem::QuantGroupSize::kN,
VecLoadSize>;
return TileEncodingPattern::make_2d_static_tile_distribution();
}

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

@@ -198,9 +198,9 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
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<sequence<MWarps, XR>,
tuple<sequence<Y0, Y1, Y2>, sequence<X>>,
@@ -213,7 +213,7 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding
{
// YPerQ > 1 implementation
// TODO: do not repeat everything to all threads
// return make_static_tile_distribution(
// tile_distribution_encoding<sequence<MWarps, 2, NWarps>,
// tuple<sequence<YPerTile, 1>, sequence<XPerTile>>,

13
test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp
View File

@@ -403,7 +403,8 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
// 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_bqn_dev_buf(bq_bqk_bqn.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
@@ -430,12 +431,12 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
// Create args for kernel execution
ck_tile::QuantGemmHostArgs args{
a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr
b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr
c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr
nullptr, // aq_ptr (not used for BQuant)
a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr
b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr
c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr
nullptr, // aq_ptr (not used for BQuant)
bq_bqk_bqn_dev_buf.GetDeviceBuffer(), // bq_ptr (scales)
1, // k_batch
1, // k_batch
M,
N,
K, // M, N, K