mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 00:58:44 +00:00
[rocm-libraries] ROCm/rocm-libraries#7199 (commit 23f7320)
[CK_TILE] [QuantGEMM] Fix SplitK tail handling and other improvements (#7199) This pull request introduces improved and more robust split-K support for quantized GEMM. The main changes add runtime validation, utility functions for split-K batch calculations, pointer offset handling for split-K in grouped kernels, and enhanced support for various tensor layouts. The changes also improve error handling and provide more flexibility for runtime tail handling in split-K pipelines. **Split-K Support and Validation Enhancements:** * Added runtime validation to ensure `k_batch` is a positive integer and that split-K configurations do not produce empty final batches or mismatched pipeline tails, with detailed error messages and logging for misconfiguration. [[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1184-R1211) [[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1161-R1250) * Introduced utility functions `get_splitk_batch_k_read` and `get_splitk_last_batch_k` to compute per-batch K read sizes and handle split rounding, ensuring correct and consistent split-K batch partitioning. [[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R206-R234) [[2]](diffhunk://#diff-635b89bdffa96b2b42f1632520cde36701d7d631e864185591f6b32f7645cf47L104-R107) [[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L388-R417) [[4]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1161-R1250) * Changed the default value of `k_batch` in `QuantGemmHostArgs` to 1 (no split-K) for safer default behavior. **Pointer Offsets and Grouped Kernel Handling:** * Updated `QuantGroupedGemmKernel` to apply split-K per-batch offsets to all input pointers, mirroring the behavior of non-grouped kernels and ensuring correctness for split-K launches. * Modified AQ tensor view handling to correctly reflect the remaining K-groups from the split-K batch's offset position, improving accuracy for split-K in grouped kernels. **Pipeline and Layout Flexibility:** * Added support for runtime selection of split-K tail handling via a new template parameter `RuntimeSplitKTail_`, with new helper methods to dispatch GEMM pipelines accordingly. [[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R273) [[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1496-R1567) [[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1427) [[4]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1447-R1629) [[5]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1459-R1641) * Improved handling for tensor layout cases, including preshuffled B and both row-major and column-major AQ layouts, ensuring correct pointer arithmetic and compatibility checks. [[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R438-R454) [[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L464-R516) [[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1184-R1211)
This commit is contained in:
Sami Remes
committed by
assistant-librarian[bot]
assistant-librarian[bot]
parent
7b9245f18c
commit
ad4e2e7624
@@ -103,9 +103,12 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str
|
||||
}();
|
||||
using BaseGemmPipeline = std::decay_t<decltype(base_gemm_pipeline)>;
|
||||
|
||||
const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
constexpr auto K1 = GemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile)
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
|
||||
@@ -180,7 +180,8 @@ struct QuantGemmHostArgs : public QuantGemmProblem
|
||||
const void* aq_ptr = nullptr;
|
||||
const void* bq_ptr = nullptr;
|
||||
void* c_ptr = nullptr;
|
||||
index_t k_batch = 0;
|
||||
// k_batch must be a positive integer; defaults to 1 (no split-K).
|
||||
index_t k_batch = 1;
|
||||
};
|
||||
|
||||
struct QuantGemmKernelArgs
|
||||
@@ -203,10 +204,35 @@ struct QuantGemmKernelArgs
|
||||
index_t k_batch;
|
||||
};
|
||||
|
||||
CK_TILE_HOST_DEVICE auto
|
||||
get_splitk_batch_k_read(index_t K, index_t k_batch, index_t k_unit) noexcept -> index_t
|
||||
{
|
||||
// k_batch and k_unit must be positive integers. Callers are expected to
|
||||
// validate via IsSupportedArgument(); this fallback returns K so a
|
||||
// misconfigured launch behaves as a no-split kernel.
|
||||
if(k_batch <= 0 || k_unit <= 0)
|
||||
{
|
||||
return K;
|
||||
}
|
||||
const index_t k_t = k_batch * k_unit;
|
||||
return (K + k_t - 1) / k_t * k_unit;
|
||||
}
|
||||
|
||||
CK_TILE_HOST_DEVICE auto
|
||||
get_splitk_last_batch_k(index_t K, index_t k_batch, index_t k_read) noexcept -> index_t
|
||||
{
|
||||
if(k_batch <= 0)
|
||||
{
|
||||
return K;
|
||||
}
|
||||
return K - k_read * (k_batch - 1);
|
||||
}
|
||||
|
||||
template <typename TilePartitioner_,
|
||||
typename GemmPipeline_,
|
||||
typename EpiloguePipeline_,
|
||||
QuantType QuantType_>
|
||||
QuantType QuantType_,
|
||||
bool RuntimeSplitKTail_ = false>
|
||||
struct QuantGemmKernel
|
||||
{
|
||||
using TilePartitioner = remove_cvref_t<TilePartitioner_>;
|
||||
@@ -244,7 +270,8 @@ struct QuantGemmKernel
|
||||
static constexpr auto I3 = number<3>(); // BQ Tensor
|
||||
static constexpr auto I4 = number<4>(); // C Tensor
|
||||
|
||||
static constexpr auto kQuantType = QuantType_;
|
||||
static constexpr auto kQuantType = QuantType_;
|
||||
static constexpr bool RuntimeSplitKTail = RuntimeSplitKTail_;
|
||||
|
||||
[[nodiscard]] CK_TILE_HOST static const std::string GetName()
|
||||
{
|
||||
@@ -386,11 +413,9 @@ struct QuantGemmKernel
|
||||
{
|
||||
constexpr auto K1 =
|
||||
GemmPipeline::BlockGemmShape::WarpTile::at(I2); // smallest unit of K work per block
|
||||
const index_t K_t = amd_wave_read_first_lane(
|
||||
kargs.k_batch * K1); // amount of K elements consumed if every split-K batch
|
||||
// performs exactly one "unit" (K1)
|
||||
const index_t KRead = amd_wave_read_first_lane(
|
||||
(kargs.K + K_t - 1) / K_t * K1); // total k elements to be read in this batch
|
||||
const index_t KRead =
|
||||
amd_wave_read_first_lane(get_splitk_batch_k_read(kargs.K, kargs.k_batch, K1));
|
||||
// total k elements to be read in this batch
|
||||
// offset not necessarily = KRead, because B can have packed elements (e.g. fp8i4)
|
||||
constexpr index_t BPackedSize =
|
||||
ck_tile::numeric_traits<remove_cvref_t<BDataType>>::PackedSize;
|
||||
@@ -412,7 +437,21 @@ struct QuantGemmKernel
|
||||
}
|
||||
else if constexpr(std::is_same_v<tensor_layout::gemm::ColumnMajor, BLayout>)
|
||||
{
|
||||
b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements);
|
||||
if constexpr(PreshuffleB)
|
||||
{
|
||||
// Preshuffled B is laid out as [N/N_Warp_Tile, K_outer, N_Warp_Tile, K_inner]
|
||||
// (see shuffle_b<>), where each "N_outer" row spans N_Warp_Tile * full_K
|
||||
// linear elements. MakeBBlockWindow already builds the descriptor with
|
||||
// stride [N_Warp_Tile * kargs.K, 1], so to advance the K starting position
|
||||
// by k_id * KRead within row 0 we need to advance the pointer by
|
||||
// (k_id * KRead) * N_Warp_Tile -- not just (k_id * KRead).
|
||||
constexpr index_t N_Warp_Tile = GemmPipeline::BlockGemmShape::WarpTile::at(I1);
|
||||
b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements * N_Warp_Tile);
|
||||
}
|
||||
else
|
||||
{
|
||||
b_k_split_offset = amd_wave_read_first_lane(b_k_offset_elements);
|
||||
}
|
||||
}
|
||||
|
||||
if(k_id < static_cast<uint32_t>(kargs.k_batch - 1))
|
||||
@@ -462,12 +501,20 @@ struct QuantGemmKernel
|
||||
using BQuantGroupSize = remove_cvref_t<typename GemmPipeline::BQuantGroupSize>;
|
||||
|
||||
// Compute AQ K-group offset for this split-K batch.
|
||||
// AQ tensor layout is RowMajor [M, QK_A] with stride [stride_AQ, 1].
|
||||
// Advancing to column aq_group_offset means a pointer offset of aq_group_offset
|
||||
// elements (column stride = 1).
|
||||
const index_t k_offset_aq = amd_wave_read_first_lane(k_id * KRead);
|
||||
aq_group_offset = amd_wave_read_first_lane(k_offset_aq / AQuantGroupSize::kK);
|
||||
aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset);
|
||||
aq_group_offset = amd_wave_read_first_lane(k_offset_aq / AQuantGroupSize::kK);
|
||||
if constexpr(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>)
|
||||
{
|
||||
// RowMajor AQ is [M, QK_A] with stride [stride_AQ, 1].
|
||||
// Advancing to K-group column g is a pointer offset of g.
|
||||
aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset);
|
||||
}
|
||||
else if constexpr(std::is_same_v<AQLayout, tensor_layout::gemm::ColumnMajor>)
|
||||
{
|
||||
// ColumnMajor AQ is [QK_A, M] with K-group row stride stride_AQ.
|
||||
// Advancing to K-group row g is a pointer offset of g * stride_AQ.
|
||||
aq_k_split_offset = amd_wave_read_first_lane(aq_group_offset * kargs.stride_AQ);
|
||||
}
|
||||
|
||||
// Compute BQ K-group offset for this split-K batch.
|
||||
// BQ tensor layout is ColumnMajor [N/kN, K/kK] with stride [K/kK, 1] for
|
||||
@@ -1135,6 +1182,34 @@ struct QuantGemmKernel
|
||||
|
||||
CK_TILE_HOST static bool IsSupportedArgument(const QuantGemmKernelArgs& kargs)
|
||||
{
|
||||
// k_batch must be a positive integer.
|
||||
if(kargs.k_batch <= 0)
|
||||
{
|
||||
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
|
||||
{
|
||||
CK_TILE_ERROR("k_batch must be a positive integer (got " +
|
||||
std::to_string(kargs.k_batch) + ")!");
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// The split-K K-unit (warp-tile K dimension) must be positive too;
|
||||
// it is a compile-time constant taken from the pipeline shape.
|
||||
static_assert(GemmPipeline::BlockGemmShape::WarpTile::at(I2) > 0,
|
||||
"Pipeline warp-tile K dimension (k_unit) must be positive.");
|
||||
|
||||
// ABQuantGrouped does not currently support RowMajor BQ layout: the
|
||||
// BQ tensor view, tile window, and split-K offset code are all
|
||||
// written for ColumnMajor BQ. The deeper static_asserts in
|
||||
// MakeBQBlockWindow enforce this at instantiation time; surface it
|
||||
// here at the host-arg entry point too so the limitation is visible
|
||||
// before the first device-side instantiation.
|
||||
static_assert(!(kQuantType == QuantType::ABQuantGrouped &&
|
||||
std::is_same_v<BQLayout, tensor_layout::gemm::RowMajor>),
|
||||
"ABQuantGrouped does not currently support RowMajor BQ layout. "
|
||||
"Use ColumnMajor BQ (or extend MakeBQBlockWindow and the split-K "
|
||||
"BQ offset path to handle RowMajor BQ).");
|
||||
|
||||
// Split-K is supported for BQuantGrouped (without preshuffle) and
|
||||
// ABQuantGrouped (without APreshuffleQuant) modes.
|
||||
if(kargs.k_batch != 1)
|
||||
@@ -1158,12 +1233,22 @@ struct QuantGemmKernel
|
||||
}
|
||||
else
|
||||
{
|
||||
constexpr auto K1 = GemmPipeline::BlockGemmShape::WarpTile::at(I2);
|
||||
const index_t K_t = kargs.k_batch * K1;
|
||||
const index_t KRead = (kargs.K + K_t - 1) / K_t * K1; // per-batch K read size
|
||||
constexpr auto K1 = GemmPipeline::BlockGemmShape::WarpTile::at(I2);
|
||||
const index_t KRead =
|
||||
get_splitk_batch_k_read(kargs.K, kargs.k_batch, K1); // per-batch K read size
|
||||
const index_t KLast = get_splitk_last_batch_k(kargs.K, kargs.k_batch, KRead);
|
||||
constexpr index_t BPackedSize =
|
||||
ck_tile::numeric_traits<remove_cvref_t<BDataType>>::PackedSize;
|
||||
|
||||
if(KLast <= 0)
|
||||
{
|
||||
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
|
||||
{
|
||||
CK_TILE_ERROR("Split-K configuration produces an empty final K batch!");
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// Constraint 1: KRead must align with B packing requirements.
|
||||
// For packed data types, multiple K elements are stored in each storage unit.
|
||||
// Split-K advances the B pointer by (KRead / BPackedSize) storage units per batch.
|
||||
@@ -1223,8 +1308,7 @@ struct QuantGemmKernel
|
||||
// (i.e. per_batch_num_loop == 1) the prefetch would read the tile
|
||||
// belonging to the next split-K batch, producing incorrect results.
|
||||
{
|
||||
const index_t per_batch_num_loop =
|
||||
KRead / static_cast<index_t>(TilePartitioner::KPerBlock);
|
||||
const index_t per_batch_num_loop = TilePartitioner::GetLoopNum(KRead);
|
||||
if(per_batch_num_loop < 2)
|
||||
{
|
||||
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
|
||||
@@ -1240,6 +1324,33 @@ struct QuantGemmKernel
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Host-side fixed tail selection is only valid when all split-K batches have
|
||||
// the same hot-loop/tail classification. Earlier batches use KRead; the final
|
||||
// batch may be shorter due to split rounding.
|
||||
{
|
||||
const index_t first_num_loop = TilePartitioner::GetLoopNum(KRead);
|
||||
const index_t last_num_loop = TilePartitioner::GetLoopNum(KLast);
|
||||
const bool first_hot_loop = GemmPipeline::BlockHasHotloop(first_num_loop);
|
||||
const bool last_hot_loop = GemmPipeline::BlockHasHotloop(last_num_loop);
|
||||
const auto first_tail = GemmPipeline::GetBlockLoopTailNum(first_num_loop);
|
||||
const auto last_tail = GemmPipeline::GetBlockLoopTailNum(last_num_loop);
|
||||
|
||||
if constexpr(!RuntimeSplitKTail)
|
||||
{
|
||||
if(first_hot_loop != last_hot_loop || first_tail != last_tail)
|
||||
{
|
||||
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
|
||||
{
|
||||
CK_TILE_ERROR(
|
||||
"Split-K batches require different hot-loop/tail handling. "
|
||||
"Use a K/k_batch combination that gives matching pipeline "
|
||||
"tails or enable runtime split-K tail dispatch.");
|
||||
}
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1383,6 +1494,78 @@ struct QuantGemmKernel
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename ADramBlockWindow, typename BDramBlockWindow, typename BQDramBlockWindow>
|
||||
CK_TILE_DEVICE static auto CallBQuantGemmPipeline(const ADramBlockWindow& a_block_window,
|
||||
const BDramBlockWindow& b_block_window,
|
||||
const BQDramBlockWindow& bq_block_window,
|
||||
const index_t num_loop,
|
||||
void* smem_ptr,
|
||||
const index_t n)
|
||||
{
|
||||
if constexpr(RuntimeSplitKTail)
|
||||
{
|
||||
static_assert(!PreshuffleB,
|
||||
"RuntimeSplitKTail is not implemented for preshuffle-B BQuant "
|
||||
"pipelines.");
|
||||
const bool has_hot_loop = GemmPipeline::BlockHasHotloop(num_loop);
|
||||
const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
return GemmPipeline{}(a_block_window,
|
||||
b_block_window,
|
||||
bq_block_window,
|
||||
num_loop,
|
||||
has_hot_loop,
|
||||
tail_num,
|
||||
smem_ptr,
|
||||
n);
|
||||
}
|
||||
else
|
||||
{
|
||||
return GemmPipeline{}(
|
||||
a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr, n);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ADramBlockWindow,
|
||||
typename BDramBlockWindow,
|
||||
typename AQDramBlockWindow,
|
||||
typename BQDramBlockWindow>
|
||||
CK_TILE_DEVICE static auto CallABQuantGemmPipeline(const ADramBlockWindow& a_block_window,
|
||||
const BDramBlockWindow& b_block_window,
|
||||
const AQDramBlockWindow& aq_block_window,
|
||||
const BQDramBlockWindow& bq_block_window,
|
||||
const index_t num_loop,
|
||||
void* smem_ptr,
|
||||
const index_t m,
|
||||
const index_t n)
|
||||
{
|
||||
if constexpr(RuntimeSplitKTail)
|
||||
{
|
||||
const bool has_hot_loop = GemmPipeline::BlockHasHotloop(num_loop);
|
||||
const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
return GemmPipeline{}(a_block_window,
|
||||
b_block_window,
|
||||
aq_block_window,
|
||||
bq_block_window,
|
||||
num_loop,
|
||||
has_hot_loop,
|
||||
tail_num,
|
||||
smem_ptr,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
else
|
||||
{
|
||||
return GemmPipeline{}(a_block_window,
|
||||
b_block_window,
|
||||
aq_block_window,
|
||||
bq_block_window,
|
||||
num_loop,
|
||||
smem_ptr,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Runs single GEMM problem cooperatively by whole workgroup.
|
||||
*
|
||||
@@ -1425,7 +1608,6 @@ struct QuantGemmKernel
|
||||
|
||||
const index_t num_loop =
|
||||
amd_wave_read_first_lane(TilePartitioner::GetLoopNum(splitk_batch_offset.splitted_k));
|
||||
|
||||
// Run GEMM cooperatively by whole workgroup.
|
||||
const auto& c_block_tile = [&]() {
|
||||
if constexpr(kQuantType == QuantType::AQuantGrouped)
|
||||
@@ -1445,7 +1627,7 @@ struct QuantGemmKernel
|
||||
{
|
||||
n = kargs.N;
|
||||
}
|
||||
return GemmPipeline{}(
|
||||
return CallBQuantGemmPipeline(
|
||||
a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr, n);
|
||||
}
|
||||
else if constexpr(kQuantType == QuantType::ABQuantGrouped)
|
||||
@@ -1457,14 +1639,14 @@ struct QuantGemmKernel
|
||||
// m = kargs.M;
|
||||
n = kargs.N;
|
||||
}
|
||||
return GemmPipeline{}(a_block_window,
|
||||
b_block_window,
|
||||
aq_block_window,
|
||||
bq_block_window,
|
||||
num_loop,
|
||||
smem_ptr,
|
||||
m,
|
||||
n);
|
||||
return CallABQuantGemmPipeline(a_block_window,
|
||||
b_block_window,
|
||||
aq_block_window,
|
||||
bq_block_window,
|
||||
num_loop,
|
||||
smem_ptr,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
else if constexpr(kQuantType == QuantType::RowColQuant ||
|
||||
kQuantType == QuantType::TensorQuant)
|
||||
@@ -1557,7 +1739,6 @@ struct QuantGemmKernel
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetSmemSize()];
|
||||
assert(kargs.k_batch == 1);
|
||||
RunGemm(
|
||||
a_ptr, b_ptr, aq_ptr, bq_ptr, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n);
|
||||
}
|
||||
|
||||
@@ -314,12 +314,18 @@ struct QuantGroupedGemmKernel
|
||||
|
||||
const typename Base::SplitKBatchOffset splitk_batch_offset(kargs, block_idx_z);
|
||||
|
||||
// options
|
||||
const ADataType* a_ptr = static_cast<const ADataType*>(kargs.a_ptr);
|
||||
const BDataType* b_ptr = static_cast<const BDataType*>(kargs.b_ptr);
|
||||
const AQDataType* aq_ptr = static_cast<const AQDataType*>(kargs.aq_ptr);
|
||||
const BQDataType* bq_ptr = static_cast<const BQDataType*>(kargs.bq_ptr);
|
||||
CDataType* c_ptr = static_cast<CDataType*>(kargs.c_ptr);
|
||||
// Apply split-K per-batch offsets to every input pointer, mirroring the
|
||||
// non-grouped QuantGemmKernel::Run_ path. All offsets are 0 when
|
||||
// k_batch == 1, so this is a no-op for non-split-K launches.
|
||||
const ADataType* a_ptr =
|
||||
static_cast<const ADataType*>(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset;
|
||||
const BDataType* b_ptr =
|
||||
static_cast<const BDataType*>(kargs.b_ptr) + splitk_batch_offset.b_k_split_offset;
|
||||
const AQDataType* aq_ptr =
|
||||
static_cast<const AQDataType*>(kargs.aq_ptr) + splitk_batch_offset.aq_k_split_offset;
|
||||
const BQDataType* bq_ptr =
|
||||
static_cast<const BQDataType*>(kargs.bq_ptr) + splitk_batch_offset.bq_k_split_offset;
|
||||
CDataType* c_ptr = static_cast<CDataType*>(kargs.c_ptr);
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetSmemSize()];
|
||||
@@ -454,8 +460,11 @@ struct QuantGroupedGemmKernel
|
||||
Base::MakeABlockWindow(a_ptr, kargs, splitk_batch_offset.splitted_k, block_idx_m);
|
||||
const auto& b_block_window =
|
||||
Base::MakeBBlockWindow(b_ptr, kargs, splitk_batch_offset.splitted_k, block_idx_n);
|
||||
const auto& aq_block_window =
|
||||
Base::MakeAQBlockWindow(aq_ptr, kargs, block_idx_m, block_idx_n);
|
||||
// Pass aq_group_offset so the AQ tensor view dimension reflects the
|
||||
// remaining K-groups from this split-K batch's offset position
|
||||
// (mirrors QuantGemmKernel::RunGemm).
|
||||
const auto& aq_block_window = Base::MakeAQBlockWindow(
|
||||
aq_ptr, kargs, block_idx_m, block_idx_n, splitk_batch_offset.aq_group_offset);
|
||||
const auto& bq_block_window = Base::MakeBQBlockWindow(
|
||||
bq_ptr, kargs, splitk_batch_offset.bq_group_offset, block_idx_m, block_idx_n);
|
||||
|
||||
|
||||
@@ -637,8 +637,10 @@ struct WPABQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRe
|
||||
const AQDramBlockWindowTmp& aq_dram_block_window_tmp,
|
||||
const BQDramBlockWindowTmp& bq_dram_block_window_tmp,
|
||||
index_t num_loop,
|
||||
bool has_hot_loop,
|
||||
TailNumber tail_number,
|
||||
void* p_smem,
|
||||
index_t m = 0,
|
||||
index_t n = 0) const
|
||||
{
|
||||
const auto RunPipeline = [&](auto bool_val, auto tail_num_) {
|
||||
@@ -650,11 +652,14 @@ struct WPABQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRe
|
||||
b_flat_dram_block_window_tmp,
|
||||
aq_dram_block_window_tmp,
|
||||
bq_dram_block_window_tmp,
|
||||
n, // dummy value, won't be used
|
||||
// The preshuffle-B ABQuant pipeline currently ignores m and n; keep this
|
||||
// runtime-tail wrapper aligned with the generic ABQuant pipeline signature.
|
||||
m,
|
||||
n,
|
||||
num_loop,
|
||||
p_smem);
|
||||
};
|
||||
return Base::TailHandler(RunPipeline, true, tail_number);
|
||||
return Base::TailHandler(RunPipeline, has_hot_loop, tail_number);
|
||||
}
|
||||
};
|
||||
} // namespace ck_tile
|
||||
|
||||
@@ -92,6 +92,11 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12")
|
||||
)
|
||||
target_compile_options(test_tile_gemm_quant_abquant_splitk_prefill PRIVATE ${TEST_GEMM_COMPILE_OPTIONS})
|
||||
|
||||
add_gtest_executable(test_tile_gemm_quant_abquant_splitk_preshuffleB
|
||||
test_gemm_quant_abquant_splitk_preshuffleB.cpp
|
||||
)
|
||||
target_compile_options(test_tile_gemm_quant_abquant_splitk_preshuffleB PRIVATE ${TEST_GEMM_COMPILE_OPTIONS})
|
||||
|
||||
add_gtest_executable(test_tile_gemm_quant_abquant_a4w4_base
|
||||
test_gemm_quant_abquant_a4w4_base.cpp
|
||||
)
|
||||
@@ -278,6 +283,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12")
|
||||
# ABQuant split-K tests
|
||||
test_tile_gemm_quant_abquant_splitk_decode
|
||||
test_tile_gemm_quant_abquant_splitk_prefill
|
||||
test_tile_gemm_quant_abquant_splitk_preshuffleB
|
||||
# BQuant tests
|
||||
test_tile_gemm_quant_bquant_1d_128
|
||||
test_tile_gemm_quant_bquant_1d_64
|
||||
|
||||
@@ -110,3 +110,46 @@ TYPED_TEST(TestCkTileGemmABQuant, SplitK8_LargeK_LargeMN)
|
||||
// K=4096, larger M and N
|
||||
this->run_test_with_validation(48, 192, 4096, 8);
|
||||
}
|
||||
|
||||
// Test one padded-K split shape whose earlier split-K batches and final batch would need
|
||||
// different compile-time pipeline tail handling. Fixed host-side tail selection rejects it.
|
||||
//
|
||||
// K=3328, k_batch=9, GemmConfigPadding (K_Tile=256):
|
||||
// K_Warp_Tile is arch-dependent (gfx94/95: 32 or 64, gfx12 WMMA: 16); for all of these
|
||||
// ceil(3328 / (9 * K_Warp_Tile)) * K_Warp_Tile = 384, so KRead is always a multiple of
|
||||
// BQuantGroupSize::kK = AQuantGroupSize::kK = 128 (Constraints 2/3 of IsSupportedArgument).
|
||||
// KLast = 3328 - 8*384 = 256.
|
||||
// num_loop_first = ceil(384/256) = 2 (hot_loop=false, tail=Even)
|
||||
// num_loop_last = ceil(256/256) = 1 (hot_loop=false, tail=Odd)
|
||||
// Mismatched tail, so the fixed host-side dispatch rejects; the runtime-tail dispatch path
|
||||
// (RuntimeSplitKTail=true) accepts and dispatches per-batch.
|
||||
using ABQuantSplitKRejectTypes = ::testing::Types<std::tuple<RowMajor,
|
||||
ColumnMajor,
|
||||
RowMajor,
|
||||
RowMajor,
|
||||
FP8,
|
||||
FP8,
|
||||
float,
|
||||
Half,
|
||||
ABQuantGrouped,
|
||||
GemmConfigPadding,
|
||||
GroupSize1x1x128,
|
||||
GroupSize1x1x128,
|
||||
ColumnMajor>>;
|
||||
|
||||
template <typename Tuple>
|
||||
class TestCkTileGemmABQuantSplitKReject : public TestCkTileGemmABQuant<Tuple>
|
||||
{
|
||||
};
|
||||
|
||||
TYPED_TEST_SUITE(TestCkTileGemmABQuantSplitKReject, ABQuantSplitKRejectTypes);
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuantSplitKReject, RejectsMismatchedTailSplitK)
|
||||
{
|
||||
EXPECT_THROW(this->run_test_with_validation(32, 128, 3328, 9), std::runtime_error);
|
||||
}
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuantSplitKReject, RuntimeTailAllowsMismatchedTailSplitK)
|
||||
{
|
||||
this->run_test_with_validation(32, 128, 3328, 9, 0, true /* allow_runtime_splitk_tail */);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include "test_gemm_quant_common.hpp"
|
||||
|
||||
using GroupSize1x1x128 = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, 128>>;
|
||||
using GroupSize1x128x128 = ck_tile::QuantGroupShape<ck_tile::sequence<1, 128, 128>>;
|
||||
|
||||
// ABQuant split-K with B-preshuffle pipeline (WPABQuantBPipelineAgBgCrV2).
|
||||
// Exercises both the regular (uniform-split) and the runtime-tail (uneven-split)
|
||||
// dispatch paths, mirroring the non-preshuffle reject/accept tests in
|
||||
// test_gemm_quant_abquant_splitk_decode.cpp.
|
||||
//
|
||||
// Tuple format: <ALayout, BLayout, CLayout, AQLayout, ADataType, BDataType, QDataType, CDataType,
|
||||
// QuantType, GemmConfig, AQuantGroupSize, BQuantGroupSize, BQLayout>
|
||||
|
||||
// =====================================================================================
|
||||
// Uniform-split tests: every split-K batch has the same per-batch num_loop and tail
|
||||
// classification, so the regular (non-runtime-tail) dispatch path applies.
|
||||
// =====================================================================================
|
||||
|
||||
// clang-format off
|
||||
using ABQuantSplitKPreshuffleBUniformTypes = ::testing::Types<
|
||||
std::tuple<RowMajor, ColumnMajor, RowMajor, RowMajor, FP8, FP8, float, Half, ABQuantGrouped, GemmConfigPreshuffleBPrefill, GroupSize1x1x128, GroupSize1x128x128, ColumnMajor>,
|
||||
std::tuple<RowMajor, ColumnMajor, RowMajor, RowMajor, BF8, BF8, float, Half, ABQuantGrouped, GemmConfigPreshuffleBPrefill, GroupSize1x1x128, GroupSize1x128x128, ColumnMajor>
|
||||
>;
|
||||
// clang-format on
|
||||
|
||||
TYPED_TEST_SUITE(TestCkTileGemmABQuant, ABQuantSplitKPreshuffleBUniformTypes);
|
||||
|
||||
// GemmConfigPreshuffleBPrefill: M_Tile=128, N_Tile=128, K_Tile=128, K_Warp_Tile=64.
|
||||
// For uniform splits we want every batch to have the same num_loop classification:
|
||||
// pick K such that KRead == KLast and num_loop_per_batch >= 2.
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK2_K1024)
|
||||
{
|
||||
// K=1024, k_batch=2 -> KRead=KLast=512, num_loop=4 per batch (Even tail).
|
||||
this->run_test_with_validation(128, 128, 1024, 2);
|
||||
}
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK4_K2048)
|
||||
{
|
||||
// K=2048, k_batch=4 -> KRead=KLast=512, num_loop=4 per batch (Even tail).
|
||||
this->run_test_with_validation(128, 128, 2048, 4);
|
||||
}
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuant, PreshuffleB_SplitK2_LargeK_LargeN)
|
||||
{
|
||||
// K=2048, larger N (multiple of N_Tile=128).
|
||||
this->run_test_with_validation(128, 256, 2048, 2);
|
||||
}
|
||||
|
||||
// =====================================================================================
|
||||
// Runtime-tail tests: K and k_batch chosen so the first split-K batch and the final
|
||||
// (shorter) batch land in different (hot-loop, tail) classifications. The default
|
||||
// host-side fixed tail dispatch must reject this; the runtime-tail dispatch path
|
||||
// (RuntimeSplitKTail=true) must accept it. Uses the padded preshuffle config so
|
||||
// uneven K passes the kPadK=false divisibility check (mirrors the non-preshuffle
|
||||
// decode test against GemmConfigPadding).
|
||||
// =====================================================================================
|
||||
|
||||
// clang-format off
|
||||
using ABQuantSplitKPreshuffleBRuntimeTailTypes = ::testing::Types<
|
||||
std::tuple<RowMajor, ColumnMajor, RowMajor, RowMajor, FP8, FP8, float, Half, ABQuantGrouped, GemmConfigPreshuffleBPrefillPadded, GroupSize1x1x128, GroupSize1x128x128, ColumnMajor>
|
||||
>;
|
||||
// clang-format on
|
||||
|
||||
template <typename Tuple>
|
||||
class TestCkTileGemmABQuantSplitKPreshuffleBReject : public TestCkTileGemmABQuant<Tuple>
|
||||
{
|
||||
};
|
||||
|
||||
TYPED_TEST_SUITE(TestCkTileGemmABQuantSplitKPreshuffleBReject,
|
||||
ABQuantSplitKPreshuffleBRuntimeTailTypes);
|
||||
|
||||
// K=3328, k_batch=9 with K_Tile=128:
|
||||
// K_Warp_Tile is arch-dependent (gfx94/95: 64 for 8-bit, gfx12 WMMA: 16); for both
|
||||
// ceil(3328 / (9 * K_Warp_Tile)) * K_Warp_Tile = 384, so KRead is always a multiple of
|
||||
// BQuantGroupSize::kK = AQuantGroupSize::kK = 128 (Constraints 2/3 of IsSupportedArgument).
|
||||
// KLast = 3328 - 8*384 = 256.
|
||||
// num_loop_first = ceil(384/128) = 3 (hot_loop=true, tail=Odd)
|
||||
// num_loop_last = ceil(256/128) = 2 (hot_loop=false, tail=Even)
|
||||
// Both hot_loop and tail differ, so the fixed host-side dispatch rejects; the runtime-tail
|
||||
// dispatch path (RuntimeSplitKTail=true) accepts and dispatches per-batch.
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuantSplitKPreshuffleBReject, RejectsMismatchedTailSplitK)
|
||||
{
|
||||
EXPECT_THROW(this->run_test_with_validation(128, 128, 3328, 9), std::runtime_error);
|
||||
}
|
||||
|
||||
TYPED_TEST(TestCkTileGemmABQuantSplitKPreshuffleBReject, RuntimeTailAllowsMismatchedTailSplitK)
|
||||
{
|
||||
this->run_test_with_validation(128, 128, 3328, 9, 0, true /* allow_runtime_splitk_tail */);
|
||||
}
|
||||
@@ -53,6 +53,31 @@ struct SafeTupleElement<TTuple,
|
||||
template <typename TTuple, size_t Index, typename DefaultType>
|
||||
using SafeTupleElement_t = typename SafeTupleElement<TTuple, Index, DefaultType>::type;
|
||||
|
||||
namespace test_gemm_quant_base_detail {
|
||||
// TODO: replace with C++20 requires later.
|
||||
// C++17 detection idiom: true when
|
||||
// T::run_quant_gemm_impl<Shape, Partitioner, Traits>(QuantGemmHostArgs, stream_config, bool)
|
||||
// is a well-formed expression.
|
||||
template <typename T, typename Shape, typename Partitioner, typename Traits, typename = void>
|
||||
struct has_run_quant_gemm_impl_splitk : std::false_type
|
||||
{
|
||||
};
|
||||
|
||||
template <typename T, typename Shape, typename Partitioner, typename Traits>
|
||||
struct has_run_quant_gemm_impl_splitk<
|
||||
T,
|
||||
Shape,
|
||||
Partitioner,
|
||||
Traits,
|
||||
std::void_t<
|
||||
decltype(std::declval<T*>()->template run_quant_gemm_impl<Shape, Partitioner, Traits>(
|
||||
std::declval<const ck_tile::QuantGemmHostArgs&>(),
|
||||
std::declval<const ck_tile::stream_config&>(),
|
||||
std::declval<bool>()))>> : std::true_type
|
||||
{
|
||||
};
|
||||
} // namespace test_gemm_quant_base_detail
|
||||
|
||||
// Base class for common quant gemm functionality
|
||||
template <typename Tuple, typename Derived>
|
||||
class TestCkTileGemmQuantBase : public ::testing::Test
|
||||
@@ -114,7 +139,9 @@ class TestCkTileGemmQuantBase : public ::testing::Test
|
||||
void TearDown() override { static_cast<Derived*>(this)->TearDownQuantTypeSpecific(); }
|
||||
|
||||
// Common test execution logic
|
||||
void invoke_quant_gemm(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s)
|
||||
void invoke_quant_gemm(const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& s,
|
||||
bool allow_runtime_splitk_tail = false)
|
||||
{
|
||||
// WP pipeline requires per-thread tile size aligned to Problem::VectorLoadSize.
|
||||
// static_assert((WG::kM * WG::kK * sizeof(ADataType) * MIterPerWarp / WaveSize) %
|
||||
@@ -149,9 +176,24 @@ class TestCkTileGemmQuantBase : public ::testing::Test
|
||||
VectorSize>;
|
||||
|
||||
// Let the derived class create the appropriate pipeline and epilogue
|
||||
static_cast<Derived*>(this)
|
||||
->template run_quant_gemm_impl<CodegenGemmShape, TilePartitioner, CodegenGemmTraits>(
|
||||
args, s);
|
||||
auto* derived = static_cast<Derived*>(this);
|
||||
if constexpr(test_gemm_quant_base_detail::has_run_quant_gemm_impl_splitk<
|
||||
Derived,
|
||||
CodegenGemmShape,
|
||||
TilePartitioner,
|
||||
CodegenGemmTraits>::value)
|
||||
{
|
||||
derived->template run_quant_gemm_impl<CodegenGemmShape,
|
||||
TilePartitioner,
|
||||
CodegenGemmTraits>(
|
||||
args, s, allow_runtime_splitk_tail);
|
||||
}
|
||||
else
|
||||
{
|
||||
derived->template run_quant_gemm_impl<CodegenGemmShape,
|
||||
TilePartitioner,
|
||||
CodegenGemmTraits>(args, s);
|
||||
}
|
||||
}
|
||||
|
||||
void RunTest(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K)
|
||||
|
||||
@@ -3,6 +3,8 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <type_traits>
|
||||
|
||||
#include "test_gemm_quant_base.hpp"
|
||||
#include "ck_tile/host/permute_pk_int4.hpp"
|
||||
#include "ck_tile/host/tensor_shuffle_utils.hpp"
|
||||
@@ -145,6 +147,12 @@ struct GemmConfigPreshuffleBPrefillTransposeC : public GemmConfigPreshuffleBPref
|
||||
static constexpr bool TransposeC = true;
|
||||
};
|
||||
|
||||
struct GemmConfigPreshuffleBPrefillPadded : public GemmConfigPreshuffleBPrefill
|
||||
{
|
||||
static constexpr bool kPadN = true;
|
||||
static constexpr bool kPadK = true;
|
||||
};
|
||||
|
||||
struct GemmConfigPreshuffleQuantPrefill : public GemmConfigPrefill
|
||||
{
|
||||
static constexpr bool BPreshuffleQuant = true;
|
||||
@@ -375,9 +383,12 @@ class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
|
||||
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
|
||||
|
||||
const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
@@ -594,16 +605,19 @@ class TestCkTileGemmAQuantMem
|
||||
void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& s)
|
||||
{
|
||||
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CodegenGemmShape,
|
||||
CodegenGemmTraits,
|
||||
ComputeDataType>;
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem<GemmPipelineProblem>;
|
||||
const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CodegenGemmShape,
|
||||
CodegenGemmTraits,
|
||||
ComputeDataType>;
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem<GemmPipelineProblem>;
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
@@ -888,9 +902,12 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGem
|
||||
ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>,
|
||||
ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2<GemmPipelineProblem>>;
|
||||
|
||||
const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
@@ -1022,8 +1039,9 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGe
|
||||
void run_test_with_validation(ck_tile::index_t M,
|
||||
ck_tile::index_t N,
|
||||
ck_tile::index_t K,
|
||||
ck_tile::index_t k_batch = 1,
|
||||
ck_tile::index_t stride_B_pad = 0)
|
||||
ck_tile::index_t k_batch = 1,
|
||||
ck_tile::index_t stride_B_pad = 0,
|
||||
bool allow_runtime_splitk_tail = false)
|
||||
{
|
||||
const ck_tile::index_t stride_A =
|
||||
ck_tile::get_default_stride(M, K, 0, this->is_row_major(ALayout{}));
|
||||
@@ -1167,7 +1185,7 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGe
|
||||
|
||||
// Run the kernel
|
||||
ck_tile::stream_config stream_config{};
|
||||
this->invoke_quant_gemm(args, stream_config);
|
||||
this->invoke_quant_gemm(args, stream_config, allow_runtime_splitk_tail);
|
||||
|
||||
// Validation using reference implementation
|
||||
ck_tile::HostTensor<CDataType> c_m_n_host_ref(
|
||||
@@ -1216,11 +1234,16 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGe
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
// ABQuant-specific pipeline implementation
|
||||
public:
|
||||
// ABQuant-specific pipeline implementation. Public so the
|
||||
// has_run_quant_gemm_impl_splitk SFINAE trait in
|
||||
// test_gemm_quant_base.hpp can detect this 3-arg overload from outside
|
||||
// the class (the trait lives in a different namespace and is not a
|
||||
// friend of this fixture).
|
||||
template <typename CodegenGemmShape, typename TilePartitioner, typename CodegenGemmTraits>
|
||||
void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args,
|
||||
const ck_tile::stream_config& s)
|
||||
const ck_tile::stream_config& s,
|
||||
bool allow_runtime_splitk_tail)
|
||||
{
|
||||
|
||||
static_assert(std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::RowMajor>);
|
||||
@@ -1258,8 +1281,10 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGe
|
||||
}();
|
||||
using BaseGemmPipeline = std::decay_t<decltype(base_gemm_pipeline)>;
|
||||
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile);
|
||||
args.k_batch == 1 ? ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile)
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
@@ -1330,23 +1355,37 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase<Tuple, TestCkTileGe
|
||||
Base::K_Warp_Tile,
|
||||
transpose_c>>>;
|
||||
|
||||
using Kernel = ck_tile::QuantGemmKernel<TilePartitioner,
|
||||
GemmPipeline,
|
||||
GemmEpilogue,
|
||||
ck_tile::QuantType::ABQuantGrouped>;
|
||||
// TODO: Replace with templated lambda when C++20 is available
|
||||
auto LaunchKernel = [&](auto RuntimeSplitKTailTag) {
|
||||
constexpr bool RuntimeSplitKTail = decltype(RuntimeSplitKTailTag)::value;
|
||||
using Kernel = ck_tile::QuantGemmKernel<TilePartitioner,
|
||||
GemmPipeline,
|
||||
GemmEpilogue,
|
||||
ck_tile::QuantType::ABQuantGrouped,
|
||||
RuntimeSplitKTail>;
|
||||
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
auto kargs = Kernel::MakeKernelArgs(args);
|
||||
const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
if(!Kernel::IsSupportedArgument(kargs))
|
||||
if(!Kernel::IsSupportedArgument(kargs))
|
||||
{
|
||||
throw std::runtime_error("Arguments not supported for ABQuant kernel");
|
||||
}
|
||||
using k_attr_t = ck_tile::kernel_attr<eight_waves>;
|
||||
ck_tile::launch_kernel(s,
|
||||
ck_tile::make_kernel<GemmConfigBase::kBlockPerCu, k_attr_t>(
|
||||
Kernel{}, grids, blocks, 0, kargs));
|
||||
};
|
||||
|
||||
if(allow_runtime_splitk_tail)
|
||||
{
|
||||
throw std::runtime_error("Arguments not supported for ABQuant kernel");
|
||||
LaunchKernel(std::true_type{});
|
||||
}
|
||||
else
|
||||
{
|
||||
LaunchKernel(std::false_type{});
|
||||
}
|
||||
using k_attr_t = ck_tile::kernel_attr<eight_waves>;
|
||||
ck_tile::launch_kernel(s,
|
||||
ck_tile::make_kernel<GemmConfigBase::kBlockPerCu, k_attr_t>(
|
||||
Kernel{}, grids, blocks, 0, kargs));
|
||||
};
|
||||
|
||||
return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
@@ -1510,9 +1549,12 @@ class TestCkTileGemmRowColQuant
|
||||
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
|
||||
|
||||
const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
@@ -1725,9 +1767,12 @@ class TestCkTileGemmTensorQuant
|
||||
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
|
||||
|
||||
const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
constexpr auto K1 = CodegenGemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split =
|
||||
args.k_batch == 1 ? (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile
|
||||
: ck_tile::get_splitk_batch_k_read(args.K, args.k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
|
||||
@@ -23,6 +23,16 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12")
|
||||
add_gtest_executable(test_ck_tile_grouped_gemm_quant_bquant_preshuffleb test_grouped_gemm_quant_bquant_preshuffleb.cpp)
|
||||
target_compile_options(test_ck_tile_grouped_gemm_quant_bquant_preshuffleb PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
|
||||
|
||||
# Smoke test for grouped ABQuant + split-K. Verifies that
|
||||
# QuantGroupedGemmKernel<..., ABQuantGrouped, ...> instantiates, that
|
||||
# IsSupportedArgument behaves correctly for valid / invalid k_batch, and
|
||||
# runs one minimal end-to-end correctness check (single group, k_batch=2)
|
||||
# to cover the grouped-specific per-batch pointer offsetting and AQ
|
||||
# group offset wiring (broader code paths are tested under
|
||||
# gemm_block_scale/test_gemm_quant_abquant_splitk_*).
|
||||
add_gtest_executable(test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke test_grouped_gemm_quant_abquant_splitk_smoke.cpp)
|
||||
target_compile_options(test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS})
|
||||
|
||||
# Collect all test targets for umbrella label
|
||||
set(CK_TILE_GROUPED_GEMM_QUANT_TEST_TARGETS
|
||||
test_ck_tile_grouped_gemm_quant_rowcol
|
||||
@@ -30,6 +40,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12")
|
||||
test_ck_tile_grouped_gemm_quant_aquant
|
||||
test_ck_tile_grouped_gemm_quant_bquant
|
||||
test_ck_tile_grouped_gemm_quant_bquant_preshuffleb
|
||||
test_ck_tile_grouped_gemm_quant_abquant_splitk_smoke
|
||||
)
|
||||
|
||||
# Label all ck_tile grouped_gemm_quant tests with CK_TILE_GROUPED_GEMM_QUANT_TESTS for selective execution
|
||||
|
||||
@@ -0,0 +1,381 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
// Minimal smoke test for QuantGroupedGemmKernel<..., ABQuantGrouped, ...>
|
||||
// with split-K. The main ABQuant + split-K code paths (uniform splits,
|
||||
// runtime-tail dispatch, BPreshuffle, etc.) are exercised in
|
||||
// test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_splitk_*.cpp using
|
||||
// the non-grouped kernel; the grouped kernel reuses the same Base::RunGemm.
|
||||
// What we cover here is grouped-specific:
|
||||
//
|
||||
// 1. Compile-time instantiation of the grouped kernel for ABQuantGrouped.
|
||||
// If any inner static_assert (e.g. RowMajor BQ + ABQuant) fires this
|
||||
// file won't compile.
|
||||
// 2. Host-side IsSupportedArgument acceptance for valid k_batch and
|
||||
// rejection for k_batch <= 0.
|
||||
// 3. A single end-to-end correctness run with k_batch == 2 on a single
|
||||
// group. This exercises QuantGroupedGemmKernel::Run's per-batch
|
||||
// pointer offsetting (a/b/aq/bq) and the aq_group_offset wired into
|
||||
// MakeAQBlockWindow -- both of which were latent bugs before this
|
||||
// change because every existing grouped test launched with k_batch=1.
|
||||
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/ops/epilogue.hpp"
|
||||
#include "ck_tile/ops/gemm.hpp"
|
||||
#include "ck_tile/ops/gemm_quant.hpp"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace {
|
||||
|
||||
using ALayout = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using BLayout = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
using CLayout = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using AQLayout = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using BQLayout = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
using ADataType = ck_tile::fp8_t;
|
||||
using BDataType = ck_tile::fp8_t;
|
||||
using AccDataType = float;
|
||||
using CDataType = ck_tile::half_t;
|
||||
using QDataType = float;
|
||||
|
||||
using ComputeDataType = ADataType;
|
||||
using AQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 1, 128>>;
|
||||
using BQuantGroupSize = ck_tile::QuantGroupShape<ck_tile::sequence<1, 128, 128>>;
|
||||
|
||||
constexpr ck_tile::index_t M_Tile = 128;
|
||||
constexpr ck_tile::index_t N_Tile = 128;
|
||||
constexpr ck_tile::index_t K_Tile = 128;
|
||||
constexpr ck_tile::index_t M_Warp = 1;
|
||||
constexpr ck_tile::index_t N_Warp = 4;
|
||||
constexpr ck_tile::index_t K_Warp = 1;
|
||||
constexpr ck_tile::index_t M_Warp_Tile = 16;
|
||||
constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
#if CK_TILE_USE_WMMA
|
||||
constexpr ck_tile::index_t K_Warp_Tile = 16;
|
||||
#else
|
||||
constexpr ck_tile::index_t K_Warp_Tile = 64;
|
||||
#endif
|
||||
constexpr bool kPadM = false;
|
||||
constexpr bool kPadN = false;
|
||||
constexpr bool kPadK = false;
|
||||
constexpr bool TransposeC = true;
|
||||
constexpr ck_tile::QuantType QuantMode = ck_tile::QuantType::ABQuantGrouped;
|
||||
|
||||
using GemmShape = ck_tile::TileGemmShape<ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
|
||||
ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
|
||||
ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
|
||||
|
||||
using TilePartitioner = ck_tile::GemmTile1DPartitioner<GemmShape>;
|
||||
|
||||
using GemmTraits = ck_tile::TileGemmQuantTraits<kPadM,
|
||||
kPadN,
|
||||
kPadK,
|
||||
/*APreshuffleQuant=*/false,
|
||||
/*BPreshuffleQuant=*/false,
|
||||
/*PreshuffleB=*/false,
|
||||
ALayout,
|
||||
BLayout,
|
||||
CLayout,
|
||||
QuantMode,
|
||||
AQLayout,
|
||||
BQLayout,
|
||||
TransposeC,
|
||||
/*DoubleSmemBuffer=*/false>;
|
||||
|
||||
// PipelineProblem template used to drive BaseGemmPipeline / TailHandler
|
||||
// dispatch. The HasHotLoop / TailNumber template parameters are fixed at
|
||||
// compile time per-instantiation; TailHandler picks the right one at
|
||||
// runtime based on (K, k_batch).
|
||||
template <bool HasHotLoop, ck_tile::TailNumber TailNum>
|
||||
using AbquantPipelineProblem =
|
||||
ck_tile::GemmABQuantPipelineProblem<ADataType,
|
||||
QDataType, // AQDataType
|
||||
BDataType,
|
||||
QDataType, // BQDataType
|
||||
AccDataType,
|
||||
GemmShape,
|
||||
GemmTraits,
|
||||
AQuantGroupSize,
|
||||
BQuantGroupSize,
|
||||
TransposeC,
|
||||
ComputeDataType,
|
||||
ck_tile::GemmPipelineScheduler::Intrawave,
|
||||
HasHotLoop,
|
||||
TailNum>;
|
||||
|
||||
using BasePipelineProblem = ck_tile::GemmPipelineProblemBase<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
GemmShape,
|
||||
GemmTraits,
|
||||
ComputeDataType>;
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<BasePipelineProblem>;
|
||||
|
||||
// Compile-time instantiation check: a representative set of types must
|
||||
// build cleanly (this is the minimal "smoke" portion of the test).
|
||||
using SmokePipeline = ck_tile::ABQuantGemmPipelineAgBgCrCompV3<
|
||||
AbquantPipelineProblem<true, ck_tile::TailNumber::Full>>;
|
||||
using SmokeEpilogue =
|
||||
ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>,
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TilePartitioner::MPerBlock,
|
||||
TilePartitioner::NPerBlock,
|
||||
M_Warp,
|
||||
N_Warp,
|
||||
M_Warp_Tile,
|
||||
N_Warp_Tile,
|
||||
K_Warp_Tile,
|
||||
TransposeC>>;
|
||||
|
||||
using SmokeGroupedKernel =
|
||||
ck_tile::QuantGroupedGemmKernel<TilePartitioner, SmokePipeline, SmokeEpilogue, QuantMode>;
|
||||
using SmokeBaseKernel = typename SmokeGroupedKernel::Base;
|
||||
|
||||
static_assert(sizeof(SmokeGroupedKernel) > 0, "QuantGroupedGemmKernel must instantiate");
|
||||
static_assert(sizeof(SmokeBaseKernel) > 0, "QuantGemmKernel base must instantiate");
|
||||
|
||||
ck_tile::QuantGemmKernelArgs MakeKargsForValidation(ck_tile::index_t k_batch)
|
||||
{
|
||||
constexpr ck_tile::index_t M = 128;
|
||||
constexpr ck_tile::index_t N = 128;
|
||||
constexpr ck_tile::index_t K = 1024;
|
||||
|
||||
ck_tile::QuantGemmKernelArgs kargs{};
|
||||
kargs.a_ptr = nullptr;
|
||||
kargs.b_ptr = nullptr;
|
||||
kargs.aq_ptr = nullptr;
|
||||
kargs.bq_ptr = nullptr;
|
||||
kargs.c_ptr = nullptr;
|
||||
kargs.M = M;
|
||||
kargs.N = N;
|
||||
kargs.K = K;
|
||||
kargs.QK_A = ck_tile::integer_divide_ceil(K, AQuantGroupSize::kK);
|
||||
kargs.QK_B = ck_tile::integer_divide_ceil(K, BQuantGroupSize::kK);
|
||||
kargs.stride_A = K;
|
||||
kargs.stride_B = K;
|
||||
kargs.stride_C = N;
|
||||
kargs.stride_AQ = kargs.QK_A;
|
||||
kargs.stride_BQ = N;
|
||||
kargs.k_batch = k_batch;
|
||||
return kargs;
|
||||
}
|
||||
|
||||
// End-to-end runner. Builds host tensors, fills them with deterministic
|
||||
// random data, runs the grouped kernel for a single group, and validates
|
||||
// against the host reference. Kept intentionally narrow: one group, fixed
|
||||
// layouts/types, no preshuffle. The wider parameter space is covered by
|
||||
// the non-grouped ABQuant tests.
|
||||
bool RunSingleGroupABQuantSplitK(ck_tile::index_t M,
|
||||
ck_tile::index_t N,
|
||||
ck_tile::index_t K,
|
||||
ck_tile::index_t k_batch)
|
||||
{
|
||||
auto is_row_major = [](auto layout) {
|
||||
return ck_tile::bool_constant<std::is_same_v<ck_tile::remove_cvref_t<decltype(layout)>,
|
||||
ck_tile::tensor_layout::gemm::RowMajor>>{};
|
||||
};
|
||||
|
||||
const ck_tile::index_t stride_A = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{}));
|
||||
const ck_tile::index_t stride_B = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{}));
|
||||
const ck_tile::index_t stride_C = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{}));
|
||||
|
||||
const ck_tile::index_t AQK = ck_tile::integer_divide_ceil(K, AQuantGroupSize::kK);
|
||||
const ck_tile::index_t BQN = ck_tile::integer_divide_ceil(N, BQuantGroupSize::kN);
|
||||
const ck_tile::index_t BQK = ck_tile::integer_divide_ceil(K, BQuantGroupSize::kK);
|
||||
const ck_tile::index_t stride_AQ =
|
||||
ck_tile::get_default_stride(M, AQK, 0, is_row_major(AQLayout{}));
|
||||
const ck_tile::index_t stride_BQ =
|
||||
ck_tile::get_default_stride(BQK, BQN, 0, is_row_major(BQLayout{}));
|
||||
|
||||
ck_tile::HostTensor<ADataType> a_m_k(
|
||||
ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(ALayout{})));
|
||||
ck_tile::HostTensor<BDataType> b_k_n(
|
||||
ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(BLayout{})));
|
||||
ck_tile::HostTensor<QDataType> aq_m_aqk(
|
||||
ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(AQLayout{})));
|
||||
ck_tile::HostTensor<QDataType> bq_bqk_bqn(
|
||||
ck_tile::host_tensor_descriptor(BQK, BQN, stride_BQ, is_row_major(BQLayout{})));
|
||||
|
||||
ck_tile::FillUniformDistribution<ADataType>{-2.0f, 3.0f}(a_m_k);
|
||||
ck_tile::FillUniformDistribution<BDataType>{-5.0f, 5.0f}(b_k_n);
|
||||
ck_tile::FillUniformDistribution<QDataType>{-2.0f, 2.0f}(aq_m_aqk);
|
||||
ck_tile::FillUniformDistribution<QDataType>{-2.0f, 2.0f}(bq_bqk_bqn);
|
||||
|
||||
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 aq_m_aqk_dev_buf(aq_m_aqk.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));
|
||||
|
||||
a_m_k_dev_buf.ToDevice(a_m_k.data());
|
||||
b_k_n_dev_buf.ToDevice(b_k_n.data());
|
||||
aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data());
|
||||
bq_bqk_bqn_dev_buf.ToDevice(bq_bqk_bqn.data());
|
||||
|
||||
if(k_batch > 1)
|
||||
{
|
||||
c_m_n_dev_buf.SetZero();
|
||||
}
|
||||
|
||||
std::vector<ck_tile::QuantGroupedGemmHostArgs> gemm_descs;
|
||||
gemm_descs.emplace_back(a_m_k_dev_buf.GetDeviceBuffer(),
|
||||
b_k_n_dev_buf.GetDeviceBuffer(),
|
||||
c_m_n_dev_buf.GetDeviceBuffer(),
|
||||
aq_m_aqk_dev_buf.GetDeviceBuffer(),
|
||||
bq_bqk_bqn_dev_buf.GetDeviceBuffer(),
|
||||
k_batch,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
AQK,
|
||||
BQK,
|
||||
stride_A,
|
||||
stride_B,
|
||||
stride_C,
|
||||
stride_AQ,
|
||||
stride_BQ);
|
||||
|
||||
// Workspace holds the per-group QuantGemmTransKernelArg vector, copied
|
||||
// to device before the launch.
|
||||
ck_tile::DeviceMem gemm_workspace(gemm_descs.size() * sizeof(ck_tile::QuantGemmTransKernelArg));
|
||||
void* kargs_ptr = gemm_workspace.GetDeviceBuffer();
|
||||
|
||||
// Drive TailHandler dispatch with split-K-aware K_split (mirrors
|
||||
// run_quant_gemm_impl in the non-grouped fixture).
|
||||
constexpr auto K1 = GemmShape::WarpTile::at(ck_tile::number<2>{});
|
||||
const ck_tile::index_t K_split = (k_batch == 1)
|
||||
? ck_tile::integer_least_multiple(K, K_Tile)
|
||||
: ck_tile::get_splitk_batch_k_read(K, k_batch, K1);
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
ck_tile::stream_config s{};
|
||||
|
||||
auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {
|
||||
constexpr bool has_hot_loop_v = has_hot_loop_.value;
|
||||
constexpr auto tail_number_v = tail_number_.value;
|
||||
|
||||
using PipelineProblem = AbquantPipelineProblem<has_hot_loop_v, tail_number_v>;
|
||||
using GemmPipeline = ck_tile::ABQuantGemmPipelineAgBgCrCompV3<PipelineProblem>;
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<
|
||||
ck_tile::CShuffleEpilogueProblem<ADataType,
|
||||
BDataType,
|
||||
ck_tile::tuple<>,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ck_tile::tuple<>,
|
||||
CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
TilePartitioner::MPerBlock,
|
||||
TilePartitioner::NPerBlock,
|
||||
M_Warp,
|
||||
N_Warp,
|
||||
M_Warp_Tile,
|
||||
N_Warp_Tile,
|
||||
K_Warp_Tile,
|
||||
TransposeC>>;
|
||||
using Kernel =
|
||||
ck_tile::QuantGroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue, QuantMode>;
|
||||
|
||||
auto kargs = Kernel::MakeKargs(gemm_descs);
|
||||
if(!Kernel::IsSupportedArgument(kargs))
|
||||
{
|
||||
throw std::runtime_error("Grouped ABQuant SplitK args not supported");
|
||||
}
|
||||
|
||||
const dim3 grids = Kernel::GridSize(gemm_descs);
|
||||
const dim3 blocks = Kernel::BlockSize();
|
||||
|
||||
HIP_CHECK_ERROR(hipMemcpyWithStream(kargs_ptr,
|
||||
kargs.data(),
|
||||
kargs.size() * sizeof(ck_tile::QuantGemmTransKernelArg),
|
||||
hipMemcpyHostToDevice,
|
||||
s.stream_id_));
|
||||
|
||||
ck_tile::launch_kernel(
|
||||
s,
|
||||
ck_tile::make_kernel<1>(Kernel{},
|
||||
grids,
|
||||
blocks,
|
||||
0,
|
||||
ck_tile::cast_pointer_to_constant_address_space(kargs_ptr),
|
||||
static_cast<ck_tile::index_t>(gemm_descs.size())));
|
||||
};
|
||||
|
||||
BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num);
|
||||
|
||||
ck_tile::HostTensor<CDataType> c_m_n_host_ref(
|
||||
ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
|
||||
c_m_n_host_ref.SetZero();
|
||||
ck_tile::reference_gemm_abquant<ADataType,
|
||||
QDataType,
|
||||
BDataType,
|
||||
QDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
AQuantGroupSize,
|
||||
BQuantGroupSize>(
|
||||
a_m_k, aq_m_aqk, b_k_n, bq_bqk_bqn, c_m_n_host_ref);
|
||||
|
||||
ck_tile::HostTensor<CDataType> c_m_n_dev_result(
|
||||
ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
|
||||
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data());
|
||||
|
||||
const float max_accumulated_value =
|
||||
*std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end());
|
||||
const auto rtol =
|
||||
std::max(ck_tile::get_relative_threshold<ADataType, CDataType, AccDataType>(
|
||||
ck_tile::integer_divide_ceil(K, k_batch)),
|
||||
ck_tile::get_relative_threshold<CDataType, CDataType, CDataType>(k_batch));
|
||||
const auto atol =
|
||||
std::max(ck_tile::get_absolute_threshold<ADataType, CDataType, AccDataType>(
|
||||
max_accumulated_value / k_batch, ck_tile::integer_divide_ceil(K, k_batch)),
|
||||
ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(
|
||||
max_accumulated_value, k_batch));
|
||||
|
||||
return ck_tile::check_err(
|
||||
c_m_n_dev_result, c_m_n_host_ref, "Grouped ABQuant SplitK mismatch", rtol, atol);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST(GroupedABQuantSplitKSmoke, AcceptsKBatch1)
|
||||
{
|
||||
EXPECT_TRUE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/1)));
|
||||
}
|
||||
|
||||
TEST(GroupedABQuantSplitKSmoke, AcceptsKBatch2)
|
||||
{
|
||||
EXPECT_TRUE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/2)));
|
||||
}
|
||||
|
||||
TEST(GroupedABQuantSplitKSmoke, RejectsKBatchZero)
|
||||
{
|
||||
EXPECT_FALSE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/0)));
|
||||
}
|
||||
|
||||
TEST(GroupedABQuantSplitKSmoke, RejectsKBatchNegative)
|
||||
{
|
||||
EXPECT_FALSE(SmokeBaseKernel::IsSupportedArgument(MakeKargsForValidation(/*k_batch=*/-1)));
|
||||
}
|
||||
|
||||
// End-to-end correctness for the grouped ABQuant + split-K path on a
|
||||
// single group with k_batch=2. Catches regressions in:
|
||||
// - QuantGroupedGemmKernel::Run per-batch a/b/aq/bq pointer offsetting,
|
||||
// - aq_group_offset wiring through MakeAQBlockWindow.
|
||||
// Both were latent (existing grouped tests only used k_batch=1).
|
||||
TEST(GroupedABQuantSplitKSmoke, EndToEnd_SingleGroup_KBatch2)
|
||||
{
|
||||
EXPECT_TRUE(RunSingleGroupABQuantSplitK(/*M=*/128,
|
||||
/*N=*/128,
|
||||
/*K=*/512,
|
||||
/*k_batch=*/2));
|
||||
}
|
||||
Reference in New Issue
Block a user