21 template <
typename,
typename Default,
typename =
void>
27 template <
typename T,
typename Default>
30 using type =
typename T::AQLayout;
33 template <
typename,
typename Default,
typename =
void>
39 template <
typename T,
typename Default>
42 using type =
typename T::BQLayout;
45 template <
typename,
typename Default,
typename =
void>
51 template <
typename T,
typename Default>
54 using type =
typename T::AQDataType;
57 template <
typename,
typename Default,
typename =
void>
63 template <
typename T,
typename Default>
66 using type =
typename T::BQDataType;
69 template <
typename,
typename =
void>
72 static constexpr
bool value =
false;
78 static constexpr
bool value = T::PreshuffleQuant;
81 template <
typename,
typename =
void>
84 static constexpr
bool value =
false;
90 static constexpr
bool value = T::PreshuffleB;
152 M_, N_, K_, QK_A_, QK_B_, stride_A_, stride_B_, stride_C_, stride_AQ_, stride_BQ_),
190 template <
typename TilePartitioner_,
191 typename GemmPipeline_,
192 typename EpiloguePipeline_,
234 return concat(
'_',
"gemm_quant", gemm_prec_str<ADataType, BDataType>, GemmPipeline::GetName());
240 return dim3(TilePartitioner::GridSize(M, N), 1, KBatch);
268 return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
274 const std::size_t k_id = blockIdx.z)
276 constexpr
auto K1 = TilePartitioner::BlockGemmShape::WarpTile::at(
I2);
280 if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
284 else if constexpr(std::is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)
289 if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, BLayout>)
293 else if constexpr(std::is_same_v<tensor_layout::gemm::ColumnMajor, BLayout>)
326 static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
327 if(kargs.
QK_A % GemmPipeline::GetVectorSizeAQ() != 0)
331 CK_TILE_ERROR(
"K_A is not a multiple of vector load size for A tensor!");
339 static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
340 if(kargs.
QK_B % GemmPipeline::GetVectorSizeBQ() != 0)
344 CK_TILE_ERROR(
"K_B is not a multiple of vector load size for B tensor!");
350 if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
352 if(kargs.
K % (TilePartitioner::KPerBlock * kargs.
k_batch) != 0 &&
353 GemmPipeline::kPadK ==
false)
357 CK_TILE_ERROR(
"Can't support K that is not a multiple of k_batch * KPerBlock "
362 if(kargs.
K % GemmPipeline::GetVectorSizeA() != 0)
366 CK_TILE_ERROR(
"K is not a multiple of vector load size for A tensor!");
373 if(kargs.
M % TilePartitioner::MPerBlock != 0 && GemmPipeline::kPadM ==
false)
378 "Can't support M that is not a multiple of MPerBlock without padding!");
382 if(kargs.
M % GemmPipeline::GetVectorSizeA() != 0)
386 CK_TILE_ERROR(
"M is not a multiple of vector load size for A tensor!");
392 if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
394 if(kargs.
N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN ==
false)
399 "Can't support N that is not a multiple of NPerBlock without padding!");
403 if(kargs.
N % GemmPipeline::GetVectorSizeB() != 0)
407 CK_TILE_ERROR(
"N is not a multiple of vector load size for B tensor!");
414 if(kargs.
K % (TilePartitioner::KPerBlock * kargs.
k_batch) != 0 &&
415 GemmPipeline::kPadK ==
false)
419 CK_TILE_ERROR(
"Can't support K that is not a multiple of k_batch * KPerBlock "
424 if(kargs.
K % GemmPipeline::GetVectorSizeB() != 0)
428 CK_TILE_ERROR(
"K is not a multiple of vector load size for B tensor!");
434 if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
436 if(kargs.
N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN ==
false)
441 "Can't support N that is not a multiple of NPerBlock without padding!");
445 if(kargs.
N % EpiloguePipeline::GetVectorSizeC() != 0)
449 CK_TILE_ERROR(
"N is not a multiple of vector load size for C tensor!");
456 if(kargs.
M % TilePartitioner::MPerBlock != 0 && GemmPipeline::kPadM ==
false)
461 "Can't support M that is not a multiple of MPerBlock without padding!");
465 if(kargs.
M % EpiloguePipeline::GetVectorSizeC() != 0)
469 CK_TILE_ERROR(
"M is not a multiple of vector load size for C tensor!");
477 template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
486 static_assert(!TilePartitioner::BlockGemmShape::PermuteA,
"Not implemented!");
487 const auto& a_tensor_view = [&]() {
488 if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
490 return make_naive_tensor_view<address_space_enum::global>(
494 number<GemmPipeline::GetVectorSizeA()>{},
499 return make_naive_tensor_view<address_space_enum::global>(
503 number<GemmPipeline::GetVectorSizeA()>{},
508 const auto get_padding_size = [](
index_t length,
index_t alignment) {
512 const auto& aq_tensor_view = [&]() {
515 static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
516 const auto aq_x = kargs.
M * GemmPipeline::KPerBlockAQ;
517 const auto aq_y = kargs.
QK_A / GemmPipeline::KPerBlockAQ;
522 number<GemmPipeline::GetVectorSizeAQ()>{},
525 const auto block_tile_size = GemmPipeline::MPerBlock * GemmPipeline::KPerBlockAQ;
534 const auto pad_aq_x = aq_pad0_desc.get_lengths()[
I1];
535 const auto wave_tile_size =
536 TilePartitioner::BlockGemmShape::WarpTile::at(
I0) * GemmPipeline::KPerBlockAQ;
537 const auto wave_tile_count_x =
548 aq_unmerge_pad0_desc,
553 wave_tile_size, get_padding_size(wave_tile_size,
get_warp_size()))),
557 const auto pad_wave_size =
566 return make_tensor_view<address_space_enum::global>(aq_ptr, aq_merge_pad1_desc);
570 static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
571 return make_naive_tensor_view<address_space_enum::global>(
575 number<GemmPipeline::GetVectorSizeAQ()>{},
580 return make_naive_tensor_view<address_space_enum::global>(
593 const auto& b_tensor_view = [&]() {
594 if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
596 if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
598 constexpr
index_t K1 = GemmPipeline::GetSmemPackB();
600 constexpr
index_t VectorSizeB =
std::min(K1, GemmPipeline::GetVectorSizeB());
601 const auto b_k0_n_k1_desc =
612 return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
616 return make_naive_tensor_view<address_space_enum::global>(
620 number<GemmPipeline::GetVectorSizeB()>{},
626 if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
628 constexpr
index_t K1 = GemmPipeline::GetSmemPackB();
630 constexpr
index_t VectorSizeB =
std::min(K1, GemmPipeline::GetVectorSizeB());
631 const auto b_k0_n_k1_desc =
642 return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
649 GemmPipeline::flatKPerWarp *
651 TilePartitioner::BlockGemmShape::WarpTile::at(
number<2>{}));
652 index_t kFlatN = kargs.
N * kargs.
K / kFlatK;
654 return make_naive_tensor_view<address_space_enum::global>(
658 number<GemmPipeline::GetVectorSizeB()>{},
663 return make_naive_tensor_view<address_space_enum::global>(
667 number<GemmPipeline::GetVectorSizeB()>{},
674 const auto& bq_tensor_view = [&]() {
677 return make_naive_tensor_view<address_space_enum::global>(
686 static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
687 return make_naive_tensor_view<address_space_enum::global>(
691 number<GemmPipeline::GetVectorSizeBQ()>{},
701 const auto& c_tensor_view = [&]() {
702 if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
704 return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
708 number<EpiloguePipeline::GetVectorSizeC()>{},
713 return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
723 a_tensor_view, aq_tensor_view, b_tensor_view, bq_tensor_view, c_tensor_view);
726 template <
typename TensorView>
729 const auto& a_pad_view = [&]() {
730 const auto& a_tensor_view = views.at(
I0);
731 if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
748 const auto& aq_pad_view = [&]() {
return views.at(
I1); }();
750 const auto& b_flat_view = views.at(
I2);
752 const auto& b_pad_view = [&]() {
753 const auto& b_tensor_view = views.at(
I2);
754 if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::ColumnMajor>)
771 const auto& bq_pad_view = [&]() {
return views.at(
I3); }();
774 const auto& c_pad_view = [&]() {
775 const auto& c_tensor_view = views.at(
I4);
776 if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
793 return make_tuple(a_pad_view, aq_pad_view, b_flat_view, bq_pad_view, c_pad_view);
797 return make_tuple(a_pad_view, aq_pad_view, b_pad_view, bq_pad_view, c_pad_view);
801 template <
typename PadView>
805 const auto& a_pad_view = views.at(
I0);
806 const auto& aq_pad_view = views.at(
I1);
807 const auto& b_pad_view = views.at(
I2);
808 const auto& bq_pad_view = views.at(
I3);
809 const auto& c_pad_view = views.at(
I4);
810 const auto& a_block_window = [&]() {
811 if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
827 const auto& aq_block_window = [&]() {
830 static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
831 constexpr
auto block_m = TilePartitioner::MPerBlock;
832 constexpr
auto warp_m = TilePartitioner::BlockGemmShape::WarpTile::at(
I0);
833 constexpr
auto aqk_per_block =
834 TilePartitioner::KPerBlock / GemmPipeline::QuantGroupSize;
835 constexpr
auto tile_window_width =
837 constexpr
auto tile_window_height = block_m / warp_m;
838 auto block_m_idx = i_m / block_m;
842 {block_m_idx * tile_window_height, 0});
846 static_assert(std::is_same_v<AQLayout, tensor_layout::gemm::RowMajor>);
847 constexpr
auto block_m = TilePartitioner::MPerBlock;
848 constexpr
auto block_k = TilePartitioner::KPerBlock;
867 const auto& b_block_window = [&]() {
874 {
static_cast<int>(i_n / TilePartitioner::BlockGemmShape::WarpTile::at(
I1)), 0});
878 if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::ColumnMajor>)
895 const auto& bq_block_window = [&]() {
905 static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>);
909 number<TilePartitioner::KPerBlock / GemmPipeline::QuantGroupSize>{}),
924 a_block_window, aq_block_window, b_block_window, bq_block_window, c_block_window);
943 template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
956 const auto& gemm_tensor_views_tuple = MakeGemmTensorViews<DstInMemOp>(
957 a_ptr, b_ptr, aq_ptr, bq_ptr, c_ptr, kargs, splitk_batch_offset);
966 const auto& a_block_window = gemm_tile_windows.at(
I0);
967 const auto& b_block_window = gemm_tile_windows.at(
I2);
969 const auto& c_block_tile = [&]() {
972 const auto& aq_block_window = gemm_tile_windows.at(
I1);
974 a_block_window, b_block_window, aq_block_window, kargs.
M, num_loop, smem_ptr_0);
978 const auto& bq_block_window = gemm_tile_windows.at(
I3);
980 a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr_0);
986 a_block_window, b_block_window, num_loop, smem_ptr_0);
991 auto& c_block_window = gemm_tile_windows.at(
I4);
996 EpiloguePipeline{}(c_block_window, c_block_tile, c_block_window, smem_ptr_0);
1000 const auto& aq_block_window = gemm_tile_windows.at(
I1);
1001 const auto& bq_block_window = gemm_tile_windows.at(
I3);
1016 const AccDataType aq_scale = type_convert<AccDataType>(*aq_ptr);
1017 const AccDataType bq_scale = type_convert<AccDataType>(*bq_ptr);
1019 c_block_window, c_block_tile, c_block_window, smem_ptr_0, aq_scale, bq_scale);
1037 template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
1051 const auto& gemm_tensor_views_tuple = MakeGemmTensorViews<DstInMemOp>(
1052 a_ptr, b_ptr, aq_ptr, bq_ptr, c_ptr, kargs, splitk_batch_offset);
1057 const index_t num_loop = __builtin_amdgcn_readfirstlane(
1058 TilePartitioner::GetLoopNum(splitk_batch_offset.
splitted_k));
1061 const auto& a_block_window = gemm_tile_windows.at(
I0);
1062 const auto& b_block_window = gemm_tile_windows.at(
I2);
1064 const auto& c_block_tile = [&]() {
1067 const auto& bq_block_window = gemm_tile_windows.at(
I3);
1082 auto& c_block_window = gemm_tile_windows.at(
I4);
1086 EpiloguePipeline{}(c_block_window, c_block_tile, c_block_window, smem_ptr_0);
1100 const auto [iM, iN] =
TilePartitioner{kargs.
M, kargs.
N}.GetOutputTileIndex(blockId);
1116 if constexpr(GemmPipeline::DoubleSmemBuffer ==
true)
1128 splitk_batch_offset,
1141 splitk_batch_offset,
#define CK_TILE_DEVICE
Definition: config.hpp:41
#define CK_TILE_HOST
Definition: config.hpp:40
#define CK_TILE_HOST_DEVICE
Definition: config.hpp:42
__host__ constexpr __device__ T min(T x)
Definition: math.hpp:116
Definition: cluster_descriptor.hpp:13
constexpr CK_TILE_HOST_DEVICE auto make_right_pad_transform(const LowLength &low_length, const RightPadLength &right_pad_, bool_constant< SkipIsValidCheck >=bool_constant< false >{})
Definition: coordinate_transform.hpp:1584
constexpr CK_TILE_HOST_DEVICE auto make_naive_tensor_descriptor(const tuple< Lengths... > &lengths, const tuple< Strides... > &strides, number< GuaranteedLastDimensionVectorLength >=number<-1 >{}, number< GuaranteedLastDimensionVectorStride >=number<-1 >{})
Definition: tensor_descriptor.hpp:274
constexpr CK_TILE_HOST_DEVICE auto integer_least_multiple(X x, Y y)
Definition: math.hpp:155
bool EnvIsEnabled(EnvVar)
Definition: env.hpp:156
constexpr CK_TILE_HOST_DEVICE auto integer_divide_ceil(X x, Y y)
Definition: math.hpp:149
void CK_TILE_ERROR(Args &&... args) noexcept
Definition: env.hpp:12
__device__ uint32_t amd_wave_read_first_lane(uint16_t v)
Definition: amd_buffer_addressing.hpp:2834
constexpr CK_TILE_HOST_DEVICE auto make_merge_transform(const LowLengths &low_lengths)
Definition: coordinate_transform.hpp:1615
int32_t index_t
Definition: integer.hpp:9
constexpr CK_TILE_HOST_DEVICE auto pad_tensor_view(const TensorView &tensor_view, const TileLengths &tile_lengths, DoPads)
Definition: tensor_view.hpp:530
constexpr CK_TILE_HOST_DEVICE auto make_pass_through_transform(const LowLength &low_length)
Definition: coordinate_transform.hpp:1558
auto concat(const Ts &... xs) -> std::enable_if_t<!AllConvertibleToStringView< Ts... >, std::string >
Definition: concat.hpp:43
remove_cv_t< std::remove_reference_t< T > > remove_cvref_t
Definition: type_traits.hpp:21
constexpr CK_TILE_HOST_DEVICE auto make_unmerge_transform(const UpLengths &up_lengths, bool_constant< Use24BitIntegerCalculation >=bool_constant< false >{})
Definition: coordinate_transform.hpp:1622
QuantType
Definition: tile_gemm_quant_traits.hpp:12
constexpr CK_TILE_HOST_DEVICE auto transform_tensor_descriptor(const OldTensorDescriptor &old_tensor_desc, const NewTransforms &new_transforms, NewLowerDimensionOldTopIdss, NewUpperDimensionNewTopIdss)
Definition: tensor_descriptor.hpp:203
constexpr CK_TILE_DEVICE auto make_tile_window(null_tensor_view, const WindowLengths &window_lengths, const multi_index< WindowLengths::size()> &, Ts &&...)
Definition: null_tile_window.hpp:75
constexpr CK_TILE_HOST_DEVICE auto make_tuple(Xs &&... xs)
Definition: tuple.hpp:360
constexpr CK_TILE_HOST_DEVICE T max(T x)
Definition: math.hpp:161
constexpr __device__ index_t get_warp_size()
Definition: get_id.hpp:10
unsigned int uint32_t
Definition: stdint.h:126
Definition: gemm_quant_kernel.hpp:133
void * c_ptr
Definition: gemm_quant_kernel.hpp:166
const void * aq_ptr
Definition: gemm_quant_kernel.hpp:164
const void * bq_ptr
Definition: gemm_quant_kernel.hpp:165
const void * b_ptr
Definition: gemm_quant_kernel.hpp:163
CK_TILE_HOST QuantGemmHostArgs()=default
index_t k_batch
Definition: gemm_quant_kernel.hpp:167
const void * a_ptr
Definition: gemm_quant_kernel.hpp:162
CK_TILE_HOST QuantGemmHostArgs(const void *a_ptr_, const void *b_ptr_, void *c_ptr_, const void *aq_ptr_, const void *bq_ptr_, index_t k_batch_, index_t M_, index_t N_, index_t K_, index_t QK_A_, index_t QK_B_, index_t stride_A_, index_t stride_B_, index_t stride_C_, index_t stride_AQ_, index_t stride_BQ_)
Definition: gemm_quant_kernel.hpp:135
Definition: gemm_quant_kernel.hpp:272
__device__ SplitKBatchOffset(const QuantGemmKernelArgs &kargs, const std::size_t k_id=blockIdx.z)
Definition: gemm_quant_kernel.hpp:273
index_t a_k_split_offset
Definition: gemm_quant_kernel.hpp:308
index_t b_k_split_offset
Definition: gemm_quant_kernel.hpp:309
index_t splitted_k
Definition: gemm_quant_kernel.hpp:310
Definition: gemm_quant_kernel.hpp:171
index_t k_batch
Definition: gemm_quant_kernel.hpp:187
index_t stride_BQ
Definition: gemm_quant_kernel.hpp:186
const void * b_ptr
Definition: gemm_quant_kernel.hpp:173
void * c_ptr
Definition: gemm_quant_kernel.hpp:176
const void * aq_ptr
Definition: gemm_quant_kernel.hpp:174
index_t stride_A
Definition: gemm_quant_kernel.hpp:182
index_t M
Definition: gemm_quant_kernel.hpp:177
const void * a_ptr
Definition: gemm_quant_kernel.hpp:172
const void * bq_ptr
Definition: gemm_quant_kernel.hpp:175
index_t QK_B
Definition: gemm_quant_kernel.hpp:181
index_t K
Definition: gemm_quant_kernel.hpp:179
index_t QK_A
Definition: gemm_quant_kernel.hpp:180
index_t stride_AQ
Definition: gemm_quant_kernel.hpp:185
index_t N
Definition: gemm_quant_kernel.hpp:178
index_t stride_C
Definition: gemm_quant_kernel.hpp:184
index_t stride_B
Definition: gemm_quant_kernel.hpp:183
Definition: gemm_quant_kernel.hpp:195
static constexpr auto I4
Definition: gemm_quant_kernel.hpp:227
static constexpr auto I3
Definition: gemm_quant_kernel.hpp:226
static constexpr bool PreshuffleB
Definition: gemm_quant_kernel.hpp:211
static constexpr CK_TILE_HOST auto GridSize(index_t M, index_t N, index_t KBatch)
Definition: gemm_quant_kernel.hpp:238
remove_cvref_t< GemmPipeline_ > GemmPipeline
Definition: gemm_quant_kernel.hpp:197
remove_cvref_t< EpiloguePipeline_ > EpiloguePipeline
Definition: gemm_quant_kernel.hpp:198
static CK_TILE_DEVICE auto MakeGemmPadViews(const TensorView &views)
Definition: gemm_quant_kernel.hpp:727
remove_cvref_t< TilePartitioner_ > TilePartitioner
Definition: gemm_quant_kernel.hpp:196
remove_cvref_t< typename EpiloguePipeline::AccDataType > AccDataType
Definition: gemm_quant_kernel.hpp:216
static constexpr auto I0
Definition: gemm_quant_kernel.hpp:223
CK_TILE_DEVICE void operator()(QuantGemmKernelArgs kargs) const
Definition: gemm_quant_kernel.hpp:1097
remove_cvref_t< typename EpiloguePipeline::ODataType > CDataType
Definition: gemm_quant_kernel.hpp:215
static constexpr index_t kBlockSize
Definition: gemm_quant_kernel.hpp:208
remove_cvref_t< typename GemmPipeline::BLayout > BLayout
Definition: gemm_quant_kernel.hpp:200
remove_cvref_t< typename GemmPipeline::CLayout > CLayout
Definition: gemm_quant_kernel.hpp:201
static CK_TILE_DEVICE auto MakeGemmTensorViews(const ADataType *a_ptr, const BDataType *b_ptr, const AQDataType *aq_ptr, const BQDataType *bq_ptr, CDataType *c_ptr, const QuantGemmKernelArgs &kargs, const SplitKBatchOffset &splitk_batch_offset)
Definition: gemm_quant_kernel.hpp:478
static constexpr auto I1
Definition: gemm_quant_kernel.hpp:224
remove_cvref_t< typename GemmPipeline::ALayout > ALayout
Definition: gemm_quant_kernel.hpp:199
static constexpr bool PreshuffleQuant
Definition: gemm_quant_kernel.hpp:209
static CK_TILE_HOST bool IsSupportedArgument(const QuantGemmKernelArgs &kargs)
Definition: gemm_quant_kernel.hpp:313
remove_cvref_t< typename detail::get_aq_data_type_or< GemmPipeline, AccDataType >::type > AQDataType
Definition: gemm_quant_kernel.hpp:219
remove_cvref_t< typename detail::get_bq_data_type_or< GemmPipeline, AccDataType >::type > BQDataType
Definition: gemm_quant_kernel.hpp:221
remove_cvref_t< typename GemmPipeline::BDataType > BDataType
Definition: gemm_quant_kernel.hpp:214
static constexpr auto I2
Definition: gemm_quant_kernel.hpp:225
static constexpr CK_TILE_HOST_DEVICE index_t GetSmemSize()
Definition: gemm_quant_kernel.hpp:266
static CK_TILE_DEVICE void RunGemm(const ADataType *a_ptr, const BDataType *b_ptr, const AQDataType *aq_ptr, const BQDataType *bq_ptr, CDataType *c_ptr, void *smem_ptr_0, const QuantGemmKernelArgs &kargs, const SplitKBatchOffset &splitk_batch_offset, const index_t block_idx_m, const index_t block_idx_n)
Runs single GEMM problem cooperatively by whole workgroup.
Definition: gemm_quant_kernel.hpp:944
static constexpr CK_TILE_HOST QuantGemmKernelArgs MakeKernelArgs(const QuantGemmHostArgs &hostArgs)
Definition: gemm_quant_kernel.hpp:246
static CK_TILE_HOST const std::string GetName()
Definition: gemm_quant_kernel.hpp:231
static CK_TILE_DEVICE auto MakeGemmTileWindows(const PadView &views, const index_t i_m, const index_t i_n)
Definition: gemm_quant_kernel.hpp:803
remove_cvref_t< typename detail::get_bq_layout_or< GemmPipeline, typename GemmPipeline::BLayout >::type > BQLayout
Definition: gemm_quant_kernel.hpp:206
static CK_TILE_DEVICE void RunGemm2LDS(const ADataType *a_ptr, const BDataType *b_ptr, const AQDataType *aq_ptr, const BQDataType *bq_ptr, CDataType *c_ptr, void *smem_ptr_0, void *smem_ptr_1, const QuantGemmKernelArgs &kargs, const SplitKBatchOffset &splitk_batch_offset, const index_t block_idx_m, const index_t block_idx_n)
Runs single GEMM problem cooperatively by whole workgroup.
Definition: gemm_quant_kernel.hpp:1038
remove_cvref_t< typename GemmPipeline::ADataType > ADataType
Definition: gemm_quant_kernel.hpp:213
remove_cvref_t< typename detail::get_aq_layout_or< GemmPipeline, typename GemmPipeline::ALayout >::type > AQLayout
Definition: gemm_quant_kernel.hpp:204
static constexpr auto kQuantType
Definition: gemm_quant_kernel.hpp:229
static constexpr CK_TILE_HOST auto BlockSize()
Definition: gemm_quant_kernel.hpp:243
Definition: gemm_quant_kernel.hpp:95
index_t stride_AQ
Definition: gemm_quant_kernel.hpp:128
index_t N
Definition: gemm_quant_kernel.hpp:121
index_t K
Definition: gemm_quant_kernel.hpp:122
index_t stride_BQ
Definition: gemm_quant_kernel.hpp:129
index_t stride_C
Definition: gemm_quant_kernel.hpp:127
index_t stride_B
Definition: gemm_quant_kernel.hpp:126
index_t stride_A
Definition: gemm_quant_kernel.hpp:125
CK_TILE_HOST QuantGemmProblem(index_t M_, index_t N_, index_t K_, index_t QK_A_, index_t QK_B_, index_t stride_A_, index_t stride_B_, index_t stride_C_, index_t stride_AQ_, index_t stride_BQ_)
Definition: gemm_quant_kernel.hpp:97
index_t QK_A
Definition: gemm_quant_kernel.hpp:123
index_t QK_B
Definition: gemm_quant_kernel.hpp:124
CK_TILE_HOST QuantGemmProblem()=default
index_t M
Definition: gemm_quant_kernel.hpp:120
Definition: integral_constant.hpp:13
typename T::AQDataType type
Definition: gemm_quant_kernel.hpp:54
Definition: gemm_quant_kernel.hpp:47
Default type
Definition: gemm_quant_kernel.hpp:48
typename T::AQLayout type
Definition: gemm_quant_kernel.hpp:30
Definition: gemm_quant_kernel.hpp:23
Default type
Definition: gemm_quant_kernel.hpp:24
typename T::BQDataType type
Definition: gemm_quant_kernel.hpp:66
Definition: gemm_quant_kernel.hpp:59
Default type
Definition: gemm_quant_kernel.hpp:60
typename T::BQLayout type
Definition: gemm_quant_kernel.hpp:42
Definition: gemm_quant_kernel.hpp:35
Default type
Definition: gemm_quant_kernel.hpp:36
Definition: gemm_quant_kernel.hpp:83
static constexpr bool value
Definition: gemm_quant_kernel.hpp:84
Definition: gemm_quant_kernel.hpp:71
static constexpr bool value
Definition: gemm_quant_kernel.hpp:72
Definition: sequence.hpp:49
#define CK_TILE_ENV(name)
Definition: env.hpp:145