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

Merge branch 'test-mixed_input_flatmm' into zhimding/ck_hot_fix_moe_sorting

Browse Source
This commit is contained in:
root
2025-09-18 09:43:02 +00:00
parent c2e0f0a0b7 31ec8bd33a
commit c3f533070a
10 changed files with 463 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
example/ck_tile/18_flatmm/flatmm_basic.cpp
View File

@@ -50,35 +50,41 @@ template <typename FlatmmConfig, typename T>
auto shuffle_b(const ck_tile::HostTensor<T>& t)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4;
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
constexpr int MaxVecSize = 16 / sizeof(T);
constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile;
constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane);
ck_tile::HostTensor<T> t_view({n_ / FlatmmConfig::N_Warp_Tile,
FlatmmConfig::N_Warp_Tile,
k_ / FlatmmConfig::K_Warp_Tile,
divisor,
FlatmmConfig::K_Warp_Tile / divisor});
k_ / ItemsPerAccess,
ItemsPerAccess});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
return ck_tile::reference_permute(t_view, {0, 2, 1, 3});
}
template <typename FlatmmConfig, typename T>
auto shuffle_b_v1(const ck_tile::HostTensor<T>& t)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4;
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
constexpr int MaxVecSize = 16 / sizeof(T);
constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile;
constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane);
constexpr int NRepeat = FlatmmConfig::N_Tile / FlatmmConfig::N_Warp_Tile / FlatmmConfig::N_Warp;
ck_tile::HostTensor<T> t_view({n_ / FlatmmConfig::N_Tile,
FlatmmConfig::N_Warp,
FlatmmConfig::N_Warp_Tile,
NRepeat,
k_ / FlatmmConfig::K_Warp_Tile,
divisor,
FlatmmConfig::K_Warp_Tile / divisor});
k_ / ItemsPerAccess,
ItemsPerAccess});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 5, 2, 6});
return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 2, 5});
}
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>

4
example/ck_tile/18_flatmm/flatmm_basic.hpp
View File

@@ -81,7 +81,7 @@ struct FlatmmConfig16
static constexpr bool DoubleSmemBuffer = false;
static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp;
static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0;
static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0;
};
template <typename DataType>
@@ -94,7 +94,7 @@ struct FlatmmConfig16_950 : public FlatmmConfig16<DataType>
static constexpr int N_Repeat =
N_Tile / FlatmmConfig16<DataType>::N_Warp_Tile / FlatmmConfig16<DataType>::N_Warp;
static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0;
static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0;
};
template <typename ADataType>

17
example/ck_tile/21_moe_flatmm/moe_flatmm.cpp
View File

@@ -35,17 +35,16 @@ auto shuffle_b(const ck_tile::HostTensor<T>& t)
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
constexpr int N_Warp_Tile = FlatmmConfig::N_Warp_Tile;
constexpr int N_Warp = FlatmmConfig::N_Warp;
constexpr int KPerLane = FlatmmConfig::K_Warp_Tile / (64 / N_Warp_Tile);
constexpr int MaxVecSize = 16 / sizeof(T);
constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile;
constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane);
ck_tile::HostTensor<T> t_view({n_ / N_Warp_Tile,
N_Warp_Tile,
k_ / (64 * KPerLane / N_Warp_Tile),
64 / N_Warp_Tile,
KPerLane});
ck_tile::HostTensor<T> t_view({n_ / FlatmmConfig::N_Warp_Tile,
FlatmmConfig::N_Warp_Tile,
k_ / ItemsPerAccess,
ItemsPerAccess});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
return ck_tile::reference_permute(t_view, {0, 2, 1, 3});
}
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>

119
include/ck_tile/core/tensor/tile_scatter_gather.hpp
View File

@@ -475,6 +475,100 @@ struct tile_scatter_gather
});
}
template <typename LdsTileWindow_,
index_t i_access_unsupport_ = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto async_load(LdsTileWindow_&& lds_tile,
number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {}) const
{
using LdsTileWindow = remove_cvref_t<LdsTileWindow_>;
using LdsDataType = typename LdsTileWindow::DataType;
using Traits = load_store_traits;
using vector_t = typename Traits::vector_t;
using SFC_Ys = typename Traits::SFC_Ys;
constexpr auto tile_dstr = TileDstr{};
// Precompute invariant values outside loops
const auto window_origin = lds_tile.get_window_origin();
const auto& bottom_tensor_view = lds_tile.get_bottom_tensor_view();
const auto& tensor_descriptor = bottom_tensor_view.get_tensor_descriptor();
auto smem_base_ptr = bottom_tensor_view.get_buffer_view().p_data_;
// loop over thread tensor space [y0, y1, ...]
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
/// TODO: use structure binding (to be captured later) if compiled in C++20
auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0];
auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1];
auto lds_window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0];
auto lds_bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1];
static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) {
constexpr auto iAccess = number<iCoord * NumAccessPerCoord + iCoordAccess>{};
// Use precomputed window origin
auto lds_bottom_tensor_thread_idx =
window_origin + lds_window_adaptor_thread_coord.get_bottom_index();
// Use precomputed tensor descriptor
const auto lds_coord =
make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx);
// Calculate SMEM address using base pointer
CK_TILE_LDS_ADDR LdsDataType* smem = smem_base_ptr + lds_coord.get_offset();
// data index [y0, y1, ...]
constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess);
constexpr auto idx_gather = idx_ys_start[number<YsGatherDim>{}];
const auto page_offset = page_idx_[idx_gather];
// merge page_offset into bottom_coord
auto mixed_bottom_thread_coord = bottom_tensor_thread_coord;
mixed_bottom_thread_coord.get_hidden_index()[number<0>{}] += page_offset;
// read from bottom tensor
if constexpr(std::is_same_v<ValidArray, std::nullptr_t>)
this->get_bottom_tensor_view().template async_get_vectorized_elements<vector_t>(
smem,
mixed_bottom_thread_coord,
number<0>{},
bool_constant<oob_conditional_check>{});
else
this->get_bottom_tensor_view().template async_get_vectorized_elements<vector_t>(
smem,
mixed_bottom_thread_coord,
number<0>{},
valids_[idx_gather],
bool_constant<oob_conditional_check>{});
// move thread coordinate
if constexpr(iCoordAccess != (NumAccessPerCoord - 1))
{
constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess);
constexpr auto forward_step_scatter = generate_tuple(
[&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; },
number<NDimY>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
forward_step_scatter);
// lds_diff doesn't need to mask the difference of the gather-dim.
constexpr auto lds_idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
idx_diff_ys);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
lds_window_adaptor_thread_coord,
lds_bottom_tensor_thread_coord,
lds_idx_diff_ps_ys);
}
});
});
}
// TODO: currently async load only implemented in inline asm
template <typename LdsTileWindow_,
index_t i_access_unsupport_ = -1,
@@ -963,6 +1057,31 @@ make_tile_scatter_gather(const TensorView_& tensor_view,
tensor_view, window_lengths, origin, tile_distribution, page_idx, valids};
}
template <typename NewTensorView_,
typename OldTensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
typename StaticPageIndexArray_,
typename StaticValidArray_,
index_t HsGatherDim = 0,
index_t NumCoord = 1>
CK_TILE_DEVICE auto replace_bottom_tensor_view(const NewTensorView_& new_tensor_view,
const tile_scatter_gather<OldTensorView_,
WindowLengths_,
StaticTileDistribution_,
StaticPageIndexArray_,
StaticValidArray_,
HsGatherDim,
NumCoord>& tile_window)
{
return make_tile_scatter_gather(new_tensor_view,
tile_window.window_lengths_,
tile_window.window_origin_,
tile_window.tile_dstr_,
tile_window.page_idx_,
tile_window.valids_);
}
template <typename TensorView,
typename WindowLengths,
typename StaticTileDistribution,

27
include/ck_tile/core/tensor/tile_window.hpp
View File

@@ -837,6 +837,24 @@ make_tile_window_raw(const TensorView_& tensor_view,
return w;
}
template <typename NewTensorView_,
typename OldTensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
index_t NumCoord = 1>
CK_TILE_DEVICE auto
replace_bottom_tensor_view(const NewTensorView_& new_tensor_view,
const tile_window_with_static_distribution<OldTensorView_,
WindowLengths_,
StaticTileDistribution_,
NumCoord>& tile_window)
{
return make_tile_window(new_tensor_view,
tile_window.get_window_lengths(),
tile_window.get_window_origin(),
tile_window.get_tile_distribution());
}
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
@@ -946,6 +964,15 @@ make_tile_window_raw(const tile_window_with_static_lengths<TensorView, WindowLen
return w;
}
template <typename NewTensorView_, typename OldTensorView_, typename WindowLengths_>
CK_TILE_DEVICE auto replace_bottom_tensor_view(
const NewTensorView_& new_tensor_view,
const tile_window_with_static_lengths<OldTensorView_, WindowLengths_>& tile_window)
{
return make_tile_window(
new_tensor_view, tile_window.get_window_lengths(), tile_window.get_window_origin());
}
template <typename TensorView_, typename WindowLengths_>
CK_TILE_DEVICE void move_tile_window(
tile_window_with_static_lengths<TensorView_, WindowLengths_>& window,

12
include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp
View File

@@ -565,13 +565,13 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1
make_tile_window(a_lds_block_ping,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
make_static_tile_distribution(typename WG::AWarpDstrEncoding{}));
PipelinePolicy::template MakeALDS_WarpTileDistribution<Problem>());
auto a_warp_window_pong_tmp =
make_tile_window(a_lds_block_pong,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
make_static_tile_distribution(typename WG::AWarpDstrEncoding{}));
PipelinePolicy::template MakeALDS_WarpTileDistribution<Problem>());
statically_indexed_array<
statically_indexed_array<decltype(a_warp_window_ping_tmp), KIterPerWarp>,
@@ -586,16 +586,10 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
a_warp_windows_ping(mIter)(kIter) = a_warp_window_ping_tmp;
a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp;
move_tile_window(a_warp_windows_ping(mIter)(kIter),
{mIter * MPerBlockPerIter, kIter * KPerBlockPerIter});
});
});
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp;
move_tile_window(a_warp_windows_pong(mIter)(kIter),
{mIter * MPerBlockPerIter, kIter * KPerBlockPerIter});
});

76
include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp
View File

@@ -250,6 +250,37 @@ struct UniversalFlatmmPipelineAgBgCrPolicy
}
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeALDS_WarpTileDistribution()
{
using TileShape = typename Problem::BlockGemmShape;
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using ALayout = remove_cvref_t<typename Problem::ALayout>;
static_assert(TileShape::BlockWarps::at(I0) == 1, "requires Wave_M == 1");
constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0);
constexpr index_t KPerXdl = Problem::BlockGemmShape::WarpTile::at(I2);
constexpr int Repeat = TileShape::BlockWarps::at(number<1>{});
constexpr int KLane = get_warp_size() / MPerXdl;
constexpr int KPerThread = KPerXdl / KLane;
constexpr int MaxVecSize = 16 / sizeof(ADataType);
constexpr int KItemsPerLoad = min(MaxVecSize, KPerThread);
constexpr int KFragment = KPerThread / KItemsPerLoad;
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<Repeat>,
tuple<sequence<MPerXdl>, sequence<KFragment, KLane, KItemsPerLoad>>,
tuple<sequence<0>, sequence<2, 1>>,
tuple<sequence<0>, sequence<1, 0>>,
sequence<2, 2>,
sequence<0, 2>>{});
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution()
{
@@ -303,10 +334,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy
{
constexpr index_t K1 = Problem::VectorLoadSize / sizeof(ADataType);
constexpr index_t K0 = KPerBlock / K1;
constexpr index_t M2 = get_warp_size() / K0;
// coalesce reading for each blocks
if constexpr(get_warp_size() % (M2 * K0) == 0)
if constexpr(get_warp_size() % K0 == 0)
{
constexpr index_t M2 = get_warp_size() / K0;
constexpr index_t M1 = BlockSize / get_warp_size();
static_assert(M2 != 0, "M2 is zero, which will lead to a division by zero error.");
static_assert(M1 != 0, "M1 is zero, which will lead to a division by zero error.");
@@ -325,18 +356,18 @@ struct UniversalFlatmmPipelineAgBgCrPolicy
}
else
{
constexpr index_t M0 = BlockSize / get_warp_size();
constexpr index_t M1 = MPerBlock / (M2 * M0);
static_assert(M0 * M1 * M2 == MPerBlock,
"Incorrect M0, M1, M2 configuration! "
"M0, M1, M2 must cover whole MPerBlock!");
constexpr index_t KWave = K0 / get_warp_size();
constexpr index_t M0 = BlockSize / get_warp_size() / KWave;
constexpr index_t M1 = MPerBlock / M0;
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<M0, M1, M2>, sequence<K0, K1>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<0>, sequence<2, 0>>,
sequence<1, 2>,
sequence<1, 1>>{});
tile_distribution_encoding<
sequence<1>,
tuple<sequence<M0, M1>, sequence<KWave, get_warp_size(), K1>>,
tuple<sequence<1, 2>, sequence<2>>,
tuple<sequence<0, 0>, sequence<1>>,
sequence<1, 2>,
sequence<1, 2>>{});
}
}
}
@@ -381,11 +412,17 @@ struct UniversalFlatmmPipelineAgBgCrPolicy
constexpr index_t WaveSize = get_warp_size();
constexpr index_t WaveNum = BlockSize / WaveSize;
constexpr index_t KBPerLoad = GetKBPerLoad<Problem>();
constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim
constexpr index_t KBPerLoad = GetKBPerLoad<Problem>();
constexpr index_t MaxVecSize = 16 / sizeof(typename Problem::BDataType);
constexpr index_t KItemsPerLoad = min(KBPerLoad, MaxVecSize);
constexpr index_t KFragment = KBPerLoad / KItemsPerLoad;
static_assert(KFragment * KItemsPerLoad == KBPerLoad);
constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim./
constexpr index_t KWavePerBlk = 1;
constexpr index_t KRepeat = 1;
static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong");
static_assert(TileShape::BlockWarps::at(number<2>{}) == 1, "Requires K_Warp == 1");
constexpr index_t NBPerLoad = 1;
constexpr index_t NThdPerWave = 1;
@@ -396,9 +433,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<WaveRepeat>, // ?
tuple<sequence<NRepeat, NWavePerBlk, NThdPerWave, NBPerLoad>, // second direction
sequence<KRepeat, KWavePerBlk, KThdPerWave, KBPerLoad>>, // first direction
sequence<WaveRepeat>, // ?
tuple<sequence<NRepeat, NWavePerBlk, NThdPerWave, NBPerLoad>, // second direction
sequence<KFragment, KWavePerBlk, KThdPerWave, KItemsPerLoad>>, // first
// direction
// wave in blk, // thd in wave
// <M, K> // <M, K>
tuple<sequence<0, 1, 2>, sequence<1, 2>>, // which direction

257
include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp
View File

@@ -119,7 +119,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
static constexpr index_t KPerBlockPerIter = kKPerBlock / KIterPerWarp;
static constexpr int MXFP4PackedSize = 2;
static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType);
static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType);
static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * MXFP4PackedSize;
static constexpr index_t m_preload = (MIterPerWarp * KIterPerWarp >= DsReadPreload)
? DsReadPreload
@@ -170,9 +170,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
static constexpr index_t ScaleBload_num =
kNPerBlock * kKPerBlock / NWarp / 32 / ScaleBload_K1 /
WaveSize; // BlockN * BlockK / NWarp / ScalePerK / ScaleB_K1 / wavesize
static constexpr index_t Bload_total_num = Bload_num_perK * KIterPerWarp + ScaleBload_num + 0X3f0;
static constexpr index_t KPerScaleLoad = KIterPerWarp / ScaleBload_num;
static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2;
static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter;
static constexpr index_t HalfMIter = (MIterPerWarp + 1) / 2;
static constexpr index_t Bload_rep = (Bload_num_perK + HalfMIter - 1) / HalfMIter;
static constexpr index_t mfma_perM_perK = NIterPerWarp * mfma_per_wg;
static constexpr index_t dswrite_mIter = (DsWritePreIssue - 1) % MIterPerWarp;
@@ -184,6 +185,11 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
CK_TILE_HOST_DEVICE static constexpr auto
SchedulerPerM(index_t dsread_perM, index_t dswrite_perM, index_t load_perM)
{
#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS
// GFX950 use BUFFER_LOAD_LDS to fill lds_buffer_A.
// There is no separate DS_WRITE instruction at all.
dswrite_perM = 0;
#endif
// Init inst order
index_t max_data_inst = dsread_perM > load_perM
? (dsread_perM > dswrite_perM ? dsread_perM : dswrite_perM)
@@ -353,32 +359,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// Calculate ds_read number per M
dsread_perM = dsread_per_wg;
// Calculate ds_write number per M
if(mIter == 0)
{
dswrite_perM =
(dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0
? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep
: 0;
}
else if(mIter >= MIterPerWarp - DsWritePreIssue + 1)
{
dswrite_perM = 0;
}
else
{
dswrite_perM = (dswrite_num_perK -
(MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0
? dswrite_rep
: 0;
}
// Add ds write when ds write data > needed
if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter))
{
if(mIter == MIterPerWarp - 1 - dswrite_mIter)
dswrite_perM = 1;
}
// Calculate buffer_load number per M
if(mIter < HalfMIter)
{
@@ -396,7 +376,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
}
if((kIter % KPerScaleLoad == 0) && (mIter == 0))
{
load_perM = load_perM + 1;
load_perM = load_perM + 1;
}
SchedulerPerM(dsread_perM, dswrite_perM, load_perM);
}
@@ -420,32 +400,6 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// Calculate ds_read number per M
dsread_perM = dsread_per_wg;
// Calculate ds_write number per M
if(mIter == 0)
{
dswrite_perM =
(dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep) > 0
? dswrite_num_perK - (MIterPerWarp - DsWritePreIssue) * dswrite_rep
: 0;
}
else if(mIter >= MIterPerWarp - DsWritePreIssue + 1)
{
dswrite_perM = 0;
}
else
{
dswrite_perM = (dswrite_num_perK -
(MIterPerWarp - DsWritePreIssue - mIter) * dswrite_rep) > 0
? dswrite_rep
: 0;
}
// Add ds write when ds write data > needed
if(dswrite_num_perK == 0 && kIter == (KIterPerWarp - 1 - dswrite_kIter))
{
if(mIter == MIterPerWarp - 1 - dswrite_mIter)
dswrite_perM = 1;
}
// Calculate buffer_load number per M
if(mIter < HalfMIter)
{
@@ -488,7 +442,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
typename AElementFunction,
typename BFlatBlockWindowTmp,
typename DequantBFlatWindow>
CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window,
CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindowTmp a_copy_dram_window_,
const AElementFunction& a_element_func,
const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp,
const DequantBFlatWindow& scale_b_flat_window,
@@ -518,37 +472,48 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
__builtin_amdgcn_sched_barrier(0);
auto a_copy_dram_window = replace_bottom_tensor_view(
PipelinePolicy::template TransformF16xF4_ATensorView<Problem>(
a_copy_dram_window_.get_bottom_tensor_view()),
a_copy_dram_window_);
// A tile in LDS
ADataType* p_a_lds_ping = static_cast<ADataType*>(p_smem_ping);
ADataType* p_a_lds_pong = static_cast<ADataType*>(p_smem_pong);
constexpr auto a_lds_block_desc =
PipelinePolicy::template MakeF16xF4_ALdsBlockDescriptor<Problem>();
constexpr auto write_a_lds_block_desc =
PipelinePolicy::template MakeF16xF4_WriteALdsBlockDescriptor<Problem>();
constexpr auto read_a_lds_block_desc =
PipelinePolicy::template MakeF16xF4_ReadALdsBlockDescriptor<Problem>();
auto a_lds_block_ping =
make_tensor_view<address_space_enum::lds>(p_a_lds_ping, a_lds_block_desc);
auto a_lds_block_pong =
make_tensor_view<address_space_enum::lds>(p_a_lds_pong, a_lds_block_desc);
auto write_a_lds_block_ping =
make_tensor_view<address_space_enum::lds>(p_a_lds_ping, write_a_lds_block_desc);
auto write_a_lds_block_pong =
make_tensor_view<address_space_enum::lds>(p_a_lds_pong, write_a_lds_block_desc);
auto read_a_lds_block_ping =
make_tensor_view<address_space_enum::lds>(p_a_lds_ping, read_a_lds_block_desc);
auto read_a_lds_block_pong =
make_tensor_view<address_space_enum::lds>(p_a_lds_pong, read_a_lds_block_desc);
auto a_copy_lds_window_ping =
make_tile_window(a_lds_block_ping,
make_tile_window(write_a_lds_block_ping,
make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
{0, 0},
PipelinePolicy::template MakeADramTileDistribution<Problem>());
auto a_copy_lds_window_pong =
make_tile_window(a_lds_block_pong,
make_tile_window(write_a_lds_block_pong,
make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
{0, 0},
PipelinePolicy::template MakeADramTileDistribution<Problem>());
// ping-pong window for A LDS
auto a_warp_window_ping_tmp =
make_tile_window(a_lds_block_ping,
make_tile_window(read_a_lds_block_ping,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution<Problem>());
auto a_warp_window_pong_tmp =
make_tile_window(a_lds_block_pong,
make_tile_window(read_a_lds_block_pong,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution<Problem>());
@@ -639,9 +604,45 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
ScaleNPerWarp>
scale_b_warp_tensor_pong;
using ABlockTile = decltype(load_tile(a_copy_dram_window));
ABlockTile a_block_tile;
enum
{
PrefillBeforeGemm = 1,
PrefillAfterGemm = 2,
PrefillAlways = PrefillBeforeGemm | PrefillAfterGemm,
};
#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS
auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a, auto prefill_location) {
// global -> lds
if constexpr(prefill_location & PrefillAfterGemm)
async_load_tile(lds_tile_a, dram_tile_a);
};
auto prefill_lds_a_stage2 = [&](auto lds_tile_a) {
// async_load_fence();
// __builtin_amdgcn_s_waitcnt(0x03fc);
// data has been stored in lds, no need more operation.
static_assert(std::is_same_v<AElementFunction, identity>,
"buffer_load_lds don't support element func fot A before mfma");
};
#else
auto prefill_lds_a_stage1 = [&](auto lds_tile_a, auto dram_tile_a, auto prefill_location) {
// global -> vgpr
if constexpr(prefill_location & PrefillBeforeGemm)
a_block_tile = load_tile(dram_tile_a);
};
auto prefill_lds_a_stage2 = [&](auto lds_tile_a) {
// vgpr -> lds
auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(lds_tile_a, a_block_tile_transformed);
};
#endif
// HEAD
// Prefetch A0
auto a_block_tile = load_tile(a_copy_dram_window);
prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window, number<PrefillAlways>{});
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
@@ -678,21 +679,19 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter});
move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp});
// A_Lds_TileDist may differ with ADramTileDistribution
auto a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(a_copy_lds_window_ping, a_block_tile_transformed);
prefill_lds_a_stage2(a_copy_lds_window_ping);
__builtin_amdgcn_sched_barrier(0);
// Prefetch A1
a_block_tile = load_tile(a_copy_dram_window);
prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window, number<PrefillAlways>{});
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
// initialize C
tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile);
__builtin_amdgcn_s_waitcnt(Bload_total_num);
block_sync_lds();
// preload A00,A10... from lds
@@ -820,14 +819,11 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
});
// Prefill A(2i+1)
a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(a_copy_lds_window_pong, a_block_tile_transformed);
prefill_lds_a_stage2(a_copy_lds_window_pong);
// Prefetch A(2i+2)
a_block_tile = load_tile(a_copy_dram_window);
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
prefill_lds_a_stage1(
a_copy_lds_window_ping, a_copy_dram_window, number<PrefillBeforeGemm>{});
// GEMM 2i
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
@@ -870,10 +866,16 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// barrier
if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last))
{
__builtin_amdgcn_s_waitcnt(Bload_total_num);
block_sync_lds();
}
});
});
prefill_lds_a_stage1(
a_copy_lds_window_ping, a_copy_dram_window, number<PrefillAfterGemm>{});
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
// move B window to next flat K
move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter});
@@ -901,8 +903,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
scale_b_flat_dram_window;
move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter),
{scale_n_iter * NFlatPerBlockPerIter,
scale_k_iter * ScaleKFlatPerWarp});
{scale_n_iter * NFlatPerBlockPerIter,
scale_k_iter * ScaleKFlatPerWarp});
scale_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) =
load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter));
@@ -915,21 +917,20 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
move_tile_window(
b_flat_dram_windows(nIter)(kIter),
{packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter +
packed_n_rank,
kIter * KFlatPerBlockPerIter});
packed_n_rank,
kIter * KFlatPerBlockPerIter});
ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter));
b_warp_tensor_ping(nIter)(kIter) = ub.u;
});
});
// Prefill A(2i+2)
a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(a_copy_lds_window_ping, a_block_tile_transformed);
prefill_lds_a_stage2(a_copy_lds_window_ping);
// Prefetch A(2i+3)
a_block_tile = load_tile(a_copy_dram_window);
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
prefill_lds_a_stage1(
a_copy_lds_window_pong, a_copy_dram_window, number<PrefillBeforeGemm>{});
// GEMM 2i+1
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
@@ -972,11 +973,16 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// barrier
if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last))
{
__builtin_amdgcn_s_waitcnt(Bload_total_num);
block_sync_lds();
}
});
});
prefill_lds_a_stage1(
a_copy_lds_window_pong, a_copy_dram_window, number<PrefillAfterGemm>{});
// move A window to next k
move_tile_window(a_copy_dram_window, {0, kKPerBlock});
// move B window to next flat K
move_tile_window(b_flat_dram_window, {0, MXFP4KPerWarp * KFlatPerBlockPerIter});
move_tile_window(scale_b_flat_dram_window, {0, ScaleKPerWarp * ScaleKFlatPerWarp});
@@ -996,7 +1002,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
if constexpr(TailNum == TailNumber::Even)
{
// prefetch B(loopK)
const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] / ContinuousKPerThread / WG::kN * ContinuousKPerThread;
const int b_k_off = b_flat_dram_window.get_tile_distribution().calculate_index()[I1] /
ContinuousKPerThread / WG::kN * ContinuousKPerThread;
static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) {
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0)
@@ -1008,8 +1015,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
scale_b_flat_dram_window;
move_tile_window(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter),
{scale_n_iter * NFlatPerBlockPerIter,
scale_k_iter * ScaleKFlatPerWarp});
{scale_n_iter * NFlatPerBlockPerIter,
scale_k_iter * ScaleKFlatPerWarp});
scale_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) =
load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter));
@@ -1017,7 +1024,8 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
});
const int b_k_off_inter = kIter * kKPerBlock / MXFP4KPerWarp + b_k_off;
if( b_k_off_inter < kKPerBlock - k_padded_zeros) {
if(b_k_off_inter < kKPerBlock - k_padded_zeros)
{
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
auto packed_n_idx = nIter / number<ContinuousScaleNPerThread>{};
auto packed_n_rank = nIter % number<ContinuousScaleNPerThread>{};
@@ -1027,18 +1035,17 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
move_tile_window(
b_flat_dram_windows(nIter)(kIter),
{packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter +
packed_n_rank,
kIter * KFlatPerBlockPerIter});
packed_n_rank,
kIter * KFlatPerBlockPerIter});
ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter));
ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter));
b_warp_tensor_pong(nIter)(kIter) = ub.u;
});
}
});
// Prefill A(loopK)
a_block_tile_transformed = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(a_copy_lds_window_pong, a_block_tile_transformed);
prefill_lds_a_stage2(a_copy_lds_window_pong);
// GEMM loopK-1
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
@@ -1082,6 +1089,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// barrier
if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last))
{
__builtin_amdgcn_s_waitcnt(Bload_total_num);
block_sync_lds();
}
});
@@ -1099,17 +1107,18 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// GEMM loopK
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
if(kIter * WG::kK < kKPerBlock - k_padded_zeros) {
if(kIter * WG::kK < kKPerBlock - k_padded_zeros)
{
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
constexpr auto AwarpIter = (kIter * MIterPerWarp + mIter) % m_preload;
static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
// read C warp tensor from C block tensor
CWarpTensor c_warp_tensor;
c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
c_warp_tensor.get_thread_buffer() =
c_block_tile.get_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
if constexpr(mIter == 0)
dequant_mxfp4(
@@ -1120,7 +1129,9 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
kIter);
// warp GEMM
WG{}(c_warp_tensor, a_warp_tensor(number<AwarpIter>{}), dequant_B_n[nIter]);
WG{}(c_warp_tensor,
a_warp_tensor(number<AwarpIter>{}),
dequant_B_n[nIter]);
// write C warp tensor into C block tensor
c_block_tile.set_y_sliced_thread_data(
@@ -1129,7 +1140,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
c_warp_tensor.get_thread_buffer());
});
if constexpr((kIter * MIterPerWarp + mIter) <
(KIterPerWarp * MIterPerWarp - m_preload))
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp;
constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp);
@@ -1189,6 +1200,7 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
// barrier
if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last))
{
__builtin_amdgcn_s_waitcnt(Bload_total_num);
block_sync_lds();
}
});
@@ -1210,15 +1222,34 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1
void* p_smem_ping,
void* p_smem_pong) const
{
return operator()(
a_dram_block_window_tmp,
[](const ADataType & a) { return a; },
b_flat_dram_block_window_tmp,
scale_b_flat_window,
num_loop,
k_padded_zeros,
p_smem_ping,
p_smem_pong);
return operator()(a_dram_block_window_tmp,
identity{},
b_flat_dram_block_window_tmp,
scale_b_flat_window,
num_loop,
k_padded_zeros,
p_smem_ping,
p_smem_pong);
}
template <typename ADramBlockWindowTmp,
typename BFlatBlockWindowTmp,
typename DequantBFlatWindow>
CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp,
const DequantBFlatWindow& scale_b_flat_window,
const index_t num_loop,
void* p_smem_ping,
void* p_smem_pong) const
{
return operator()(a_dram_block_window_tmp,
identity{},
b_flat_dram_block_window_tmp,
scale_b_flat_window,
num_loop,
0,
p_smem_ping,
p_smem_pong);
}
};

110
include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp
View File

@@ -7,6 +7,18 @@
namespace ck_tile {
#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 0
#if defined(__gfx950__)
#define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 1
#else
#define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 0
#endif
#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS \
(CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE && \
CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4)
struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
{
static constexpr auto I0 = number<0>{};
@@ -17,16 +29,72 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
static constexpr index_t N_Pack = 2; // it's fixed for fp4
static constexpr index_t K_Pack = 2; // it's fixed for fp4
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ALdsBlockDescriptor()
template <typename Problem, typename NativeADramTensorView>
CK_TILE_HOST_DEVICE static constexpr auto
TransformF16xF4_ATensorView(const NativeADramTensorView& a_dram_view)
{
using namespace ck_tile;
#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS
constexpr int DynamicTileOffsetFlag = 0;
constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0);
constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1);
static_assert(MPerXdl == 16 && NPerXdl == 16);
constexpr index_t MPerBlock = Problem::BlockGemmShape::kM;
constexpr index_t KPerBlock = Problem::BlockGemmShape::kK;
constexpr index_t KPack = GetSmemPackA<Problem>();
constexpr int ContiguousThreadsCntInDS_READ_16B = 4;
// implement swizzle pattern on global side
// because we can't adjust the ds_write pattern of BUFFER_LOAD_LDS.
auto swizzle_a_dram_view_1 = transform_tensor_view(
a_dram_view,
make_tuple(
// M-dim is not affected by swizzle pattern
make_unmerge_transform(
make_tuple(number<DynamicTileOffsetFlag>{}, number<MPerBlock>{})),
// K-dim is the swizzle dimension
make_unmerge_transform(make_tuple(number<DynamicTileOffsetFlag>{},
number<KPerBlock / KPack>{},
number<KPack>{}))),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}));
auto swizzle_a_dram_view_2 = transform_tensor_view(
swizzle_a_dram_view_1,
make_tuple(make_pass_through_transform(number<DynamicTileOffsetFlag>{}),
make_xor_transform(make_tuple(number<MPerBlock>{},
number<ContiguousThreadsCntInDS_READ_16B>{})),
make_pass_through_transform(number<DynamicTileOffsetFlag>{}),
make_pass_through_transform(number<KPack>{})),
make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{}),
make_tuple(sequence<0>{}, sequence<1, 3>{}, sequence<2>{}, sequence<4>{}));
return transform_tensor_view(
swizzle_a_dram_view_2,
make_tuple(
make_merge_transform_v3_division_mod(
make_tuple(number<DynamicTileOffsetFlag>{}, number<MPerBlock>{})),
make_merge_transform_v3_division_mod(make_tuple(number<DynamicTileOffsetFlag>{},
number<KPerBlock / KPack>{},
number<KPack>{}))),
make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
#else
return a_dram_view;
#endif
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ReadALdsBlockDescriptor()
{
constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0);
constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1);
static_assert(MPerXdl == 16 && NPerXdl == 16);
/*reduce transform layers,compare with old ck*/
constexpr index_t MPerBlock = Problem::BlockGemmShape::kM;
constexpr index_t KPerBlock = Problem::BlockGemmShape::kK;
@@ -60,35 +128,19 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeFp16xF4_ADramTileDistribution()
CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_WriteALdsBlockDescriptor()
{
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using ALayout = remove_cvref_t<typename Problem::ALayout>;
constexpr index_t BlockSize = Problem::kBlockSize;
#if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS
constexpr index_t MPerBlock = Problem::BlockGemmShape::kM;
constexpr index_t KPerBlock = Problem::BlockGemmShape::kK;
constexpr index_t K1 = Problem::VectorLoadSize / sizeof(ADataType);
constexpr index_t K0 = KPerBlock / K1;
constexpr index_t M2 = get_warp_size() / K0;
constexpr index_t M1 = BlockSize / get_warp_size();
static_assert(M2 != 0, "M2 is zero, which will lead to a division by zero error.");
static_assert(M1 != 0, "M1 is zero, which will lead to a division by zero error.");
constexpr index_t M0 = MPerBlock / (M2 * M1);
static_assert(M0 * M1 * M2 == MPerBlock,
"Incorrect M0, M2, M1 configuration! "
"M0, M1, M2 must cover whole MPerBlock!");
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<M0, M1, M2>, sequence<K0, K1>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<1>, sequence<2, 0>>,
sequence<1, 2>,
sequence<0, 1>>{});
constexpr index_t KPack = GetSmemPackA<Problem>();
return make_naive_tensor_descriptor(make_tuple(number<MPerBlock>{}, number<KPerBlock>{}),
make_tuple(number<KPerBlock>{}, number<1>{}),
number<KPack>{},
number<1>{});
#else
return MakeF16xF4_ReadALdsBlockDescriptor<Problem>();
#endif
}
template <typename Problem>

4
include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp
View File

@@ -508,13 +508,13 @@ struct MoeFlatmmPipelineAGmemBGmemCRegV1
make_tile_window(a_lds_block_ping,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
make_static_tile_distribution(typename WG::AWarpDstrEncoding{}));
PipelinePolicy::template MakeALDS_WarpTileDistribution<Problem>());
auto a_warp_window_pong_tmp =
make_tile_window(a_lds_block_pong,
make_tuple(number<WG::kM>{}, number<WG::kK>{}),
{iMWarp * WG::kM, 0},
make_static_tile_distribution(typename WG::AWarpDstrEncoding{}));
PipelinePolicy::template MakeALDS_WarpTileDistribution<Problem>());
statically_indexed_array<
statically_indexed_array<decltype(a_warp_window_ping_tmp), KIterPerWarp>,