From a1ed60c0879fdf80435de4137cd4d4a7ba3652e6 Mon Sep 17 00:00:00 2001 From: zanzhang Date: Fri, 5 Dec 2025 13:40:11 +0800 Subject: [PATCH] update --- ...ec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 1449 +++++++++-------- ...mm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 78 +- 2 files changed, 781 insertions(+), 746 deletions(-) diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 9f77890b5f..f1865410e4 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -1266,38 +1266,37 @@ template struct F8xMXF4FlatmmPipelineProblem : FlatmmPipelineProblem + ADataType_, + CDataType_, + BlockGemmShape_, + Traits_, + Scheduler_, + HasHotLoop_, + TailNum_, + ComputeDataType_> { using BlockGemmShape = BlockGemmShape_; - // using QuantType = BDataType_; + using QuantType = BDataType_; static constexpr index_t flatNPerWarp = BlockGemmShape::flatNPerWarp; - static constexpr int ScaleGranularityK = 32; + static constexpr int MXF4ScaleGranularityK = 32; - static constexpr int ContinuousKPerThread = 32; // it's fixed for fp4 - static constexpr int MXdlPack = 2; // it's fixed for fp4 - static constexpr int NXdlPack = 2; // it's fixed for fp4 - static constexpr int KXdlPack = 2; - // static constexpr index_t flatKPerWarp = BlockGemmShape::flatKPerWarp * KXdlPack; - static constexpr index_t flatKPerWarp = get_warp_size() * ContinuousKPerThread; + static constexpr int ContinuousKPerThread = 32; // it's fixed for fp4 + static constexpr int ContinuousScaleNPerThread = 2; // it's fixed for fp4 + static constexpr int ContinuousScaleKPerThread = 2; // it's fixed for fp4 + static constexpr index_t flatKPerWarp = 64 * ContinuousKPerThread; }; template -struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 +struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 + : FlatmmPipelineAGmemBGmemCRegV1 { using Underlying = FlatmmPipelineAGmemBGmemCRegV1; using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape @@ -1317,7 +1316,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1())>; static constexpr index_t DsWritePreIssue = 3; // default 2, ds write at MIter - 2 - static constexpr index_t DsReadPreload = 4; // default 4 for MXFP4 (MXdlPack * KXdlPack) + static constexpr index_t DsReadPreload = 2; // default 2, preload 2 ds read static constexpr index_t BlockSize = Problem::kBlockSize; static constexpr index_t WaveSize = get_warp_size(); @@ -1329,17 +1328,17 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1(); static constexpr auto I1 = number<1>(); @@ -1364,45 +1363,63 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1::PackedSize; - static constexpr index_t BPackedSize = numeric_traits::PackedSize; - - static constexpr index_t MXdlPack = Problem::MXdlPack; - static constexpr index_t NXdlPack = Problem::NXdlPack; - static constexpr index_t KXdlPack = Problem::KXdlPack; - static constexpr index_t ScaleGranularityK = Problem::ScaleGranularityK; - - static constexpr index_t AK1 = Problem::VectorLoadSize / sizeof(ADataType) * APackedSize; - static constexpr index_t BK1 = Problem::VectorLoadSize / sizeof(BDataType) * BPackedSize; - + static constexpr int MXFP4PackedSize = 2; + 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 : MIterPerWarp * KIterPerWarp; + static constexpr int ContinuousKPerThread = Problem::ContinuousKPerThread; + static constexpr int ContinuousScaleNPerThread = Problem::ContinuousScaleNPerThread; + static constexpr int ContinuousScaleKPerThread = Problem::ContinuousScaleKPerThread; + + static constexpr int ScaleKFlatPerWarp = + ContinuousScaleNPerThread * ContinuousScaleKPerThread * get_warp_size(); + + static constexpr int XDLK_PerThread = + WarpTile::at(I2) / (get_warp_size() / WarpTile::at(I1)); // 8 + + static constexpr int XDL_PerWeightK = 4; // 4 + static constexpr int XDL_PerScaleK = XDL_PerWeightK * ContinuousScaleKPerThread; // 4 + static constexpr int XDL_PerScaleN = ContinuousScaleNPerThread; // 2 + static_assert(XDL_PerScaleK % XDL_PerWeightK == 0); + static_assert(KIterPerWarp % XDL_PerScaleK == 0); + static_assert(NIterPerWarp % XDL_PerScaleN == 0); + + static constexpr int MXFP4KPerWarp = KIterPerWarp / XDL_PerWeightK; + static constexpr int ScaleKPerWarp = KIterPerWarp / XDL_PerScaleK; + static constexpr int ScaleNPerWarp = NIterPerWarp / XDL_PerScaleN; + + static constexpr int MXFP4K_PerScaleK = MXFP4KPerWarp / ScaleKPerWarp; + static constexpr bool HasHotLoop = Problem::HasHotLoop; static constexpr auto TailNum = Problem::TailNum; - static constexpr index_t mfma_per_wg = 1; // 950 only - - static constexpr index_t dsread_per_wg = WG::kM * WG::kK / AK1 / WaveSize; - static_assert((WG::kM * WG::kK) % (AK1 * WaveSize) == 0); +#ifdef __gfx942__ + static constexpr index_t mfma_per_wg = 2; +#else + static constexpr index_t mfma_per_wg = 1; +#endif + static constexpr index_t dsread_per_wg = + WG::kM * WG::kK * sizeof(ADataType) / WaveSize / Problem::VectorLoadSize; + static_assert((WG::kM * WG::kK * sizeof(ADataType) / WaveSize) % Problem::VectorLoadSize == 0); static constexpr index_t dsread_num_perK = dsread_per_wg * MIterPerWarp; - static constexpr index_t dswrite_num_perK = dsread_num_perK / NWarp; + static constexpr index_t dswrite_num_perK = dsread_num_perK / (MWarp * NWarp); static constexpr index_t dswrite_rep = (dswrite_num_perK + MIterPerWarp - 1) / MIterPerWarp; static constexpr index_t Aload_num_perK = dswrite_num_perK; static constexpr index_t Aload_rep = dswrite_rep; - static constexpr index_t Bload_num_perK = kNPerBlock * WG::kK / NWarp / BK1 / WaveSize; - static constexpr index_t Bload_num = Bload_num_perK * KIterPerWarp; + static constexpr index_t ScaleBload_K1 = ContinuousScaleNPerThread * ContinuousScaleKPerThread; static constexpr index_t ScaleBload_num = - kNPerBlock * kKPerBlock / NWarp / ScaleGranularityK / NXdlPack / KXdlPack / WaveSize; - static constexpr index_t ScaleAload_num = - kMPerBlock * kKPerBlock / MWarp / ScaleGranularityK / MXdlPack / KXdlPack / WaveSize; - - // 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; + 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 mfma_perM_perK = NIterPerWarp * mfma_per_wg; static constexpr index_t dswrite_mIter = (DsWritePreIssue - 1) % MIterPerWarp; @@ -1414,6 +1431,11 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 load_perM ? (dsread_perM > dswrite_perM ? dsread_perM : dswrite_perM) @@ -1465,7 +1487,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 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) { @@ -1625,10 +1620,10 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 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) { @@ -1716,20 +1685,18 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1 - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_copy_dram_window_tmp, - const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, - const ScaleADramBlockWindowTmp& scale_a_window, - const ScaleBDramBlockWindowTmp& scale_b_window, - index_t num_loop, - void* __restrict__ p_smem_ping, - void* __restrict__ p_smem_pong) const + typename DequantBFlatWindow> + 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, + const index_t num_loop, + const index_t k_padded_zeros, + void* p_smem_ping, + void* p_smem_pong) const { -#ifndef __gfx950__ - static_assert(false, "Only gfx950 is supported for MXFP4 flatmm pipeline now."); -#endif static_assert( std::is_same_v>, "wrong!"); @@ -1739,8 +1706,8 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}], "wrong!"); - // constexpr auto MIter_2nd_last = max(0, MIterPerWarp - 2); - static_assert(NWarp == 4); + constexpr auto MIter_2nd_last = (MIterPerWarp >= 2) ? MIterPerWarp - 2 : MIterPerWarp - 1; + const index_t iMWarp = get_warp_id() / NWarp; using CWarpDstr = typename WG::CWarpDstr; using CWarpTensor = typename WG::CWarpTensor; @@ -1749,43 +1716,82 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1{}; - auto a_dram_window = - make_tile_window(PipelinePolicy::template MakeMXFP4_AAsyncLoadDramDescriptor( - a_copy_dram_window_tmp.get_bottom_tensor_view()), - a_copy_dram_window_tmp.get_window_lengths(), - a_copy_dram_window_tmp.get_window_origin(), - PipelinePolicy::template MakeMXFP4_ADramTileDistribution()); - __builtin_amdgcn_sched_barrier(0); + auto a_copy_dram_window = replace_bottom_tensor_view( + PipelinePolicy::template TransformF16xF4_ATensorView( + a_copy_dram_window_.get_bottom_tensor_view()), + a_copy_dram_window_); + // A tile in LDS ADataType* p_a_lds_ping = static_cast(p_smem_ping); ADataType* p_a_lds_pong = static_cast(p_smem_pong); - constexpr auto a_lds_block_desc = - PipelinePolicy::template MakeMXFP4_ALdsBlockDescriptor(); + constexpr auto write_a_lds_block_desc = + PipelinePolicy::template MakeF16xF4_WriteALdsBlockDescriptor(); + constexpr auto read_a_lds_block_desc = + PipelinePolicy::template MakeF16xF4_ReadALdsBlockDescriptor(); - auto a_lds_block_ping = - make_tensor_view(p_a_lds_ping, a_lds_block_desc); - auto a_lds_block_pong = - make_tensor_view(p_a_lds_pong, a_lds_block_desc); + auto write_a_lds_block_ping = + make_tensor_view(p_a_lds_ping, write_a_lds_block_desc); + auto write_a_lds_block_pong = + make_tensor_view(p_a_lds_pong, write_a_lds_block_desc); + auto read_a_lds_block_ping = + make_tensor_view(p_a_lds_ping, read_a_lds_block_desc); + auto read_a_lds_block_pong = + make_tensor_view(p_a_lds_pong, read_a_lds_block_desc); - auto a_store_lds_window_ping = make_tile_window( - a_lds_block_ping, make_tuple(number{}, number{}), {0, 0}); - auto a_store_lds_window_pong = make_tile_window( - a_lds_block_pong, make_tuple(number{}, number{}), {0, 0}); + auto a_copy_lds_window_ping = + make_tile_window(write_a_lds_block_ping, + make_tuple(number{}, number{}), + {0, 0}, + PipelinePolicy::template MakeADramTileDistribution()); + auto a_copy_lds_window_pong = + make_tile_window(write_a_lds_block_pong, + make_tuple(number{}, number{}), + {0, 0}, + PipelinePolicy::template MakeADramTileDistribution()); // ping-pong window for A LDS - auto a_warp_window_ping = - make_tile_window(a_lds_block_ping, + auto a_warp_window_ping_tmp = + make_tile_window(read_a_lds_block_ping, make_tuple(number{}, number{}), - {0, 0}, - PipelinePolicy::template MakeF8xMXF4_ALDS_TileDistribution()); - auto a_warp_window_pong = - make_tile_window(a_lds_block_pong, + {iMWarp * WG::kM, 0}, + PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution()); + auto a_warp_window_pong_tmp = + make_tile_window(read_a_lds_block_pong, make_tuple(number{}, number{}), - {0, 0}, - PipelinePolicy::template MakeF8xMXF4_ALDS_TileDistribution()); + {iMWarp * WG::kM, 0}, + PipelinePolicy::template MakeF16xF4_ALDS_TileDistribution()); + + statically_indexed_array< + statically_indexed_array, + MIterPerWarp> + a_warp_windows_ping; + + statically_indexed_array< + statically_indexed_array, + MIterPerWarp> + a_warp_windows_pong; + + auto A_Lds_Stride = 8; + 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; + + auto weight_k_idx = kIter / number{}; + auto weight_k_rank = kIter % number{}; + move_tile_window( + a_warp_windows_ping(mIter)(kIter), + {mIter * MPerBlockPerIter, + weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); + move_tile_window( + a_warp_windows_pong(mIter)(kIter), + {mIter * MPerBlockPerIter, + weight_k_rank * A_Lds_Stride + weight_k_idx * XDL_PerWeightK * WG::kK}); + }); + }); // Block GEMM auto block_flatmm = BlockFlatmm(); @@ -1793,611 +1799,706 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1(); + auto scale_b_flat_distribution = + PipelinePolicy::template MakeFp4ScaleBFlatDramTileDistribution(); + + auto b_flat_dram_window = make_tile_window( + b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views + make_tuple(number{}, number{}), + b_flat_dram_block_window_tmp.get_window_origin(), + b_flat_distribution); + + auto scale_b_flat_dram_window = make_tile_window( + scale_b_flat_window.get_bottom_tensor_view(), // from kernel gemm_pad_views + make_tuple(number{}, number{}), + scale_b_flat_window.get_window_origin(), + scale_b_flat_distribution); + + using MXFP4_Buffer = decltype(load_tile(b_flat_dram_window)); + // use v4i32 as the data type between basicblock to avoid unpack and repack operation. + using V4UInt_Buffer = thread_buffer; + union UnionB + { + V4UInt_Buffer u = 0; + MXFP4_Buffer mxfp4; + } ub; // pingpong buffer for B - auto b_flat_dram_windows = generate_tuple( - [&](auto nIter) { - constexpr auto packed_n_idx = nIter / number{}; - constexpr auto packed_n_rank = nIter % number{}; - auto window_i = make_tile_window( - b_flat_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_flat_dram_block_window_tmp.get_window_origin(), - PipelinePolicy::template MakeMXFP4_BFlatDramTileDistribution()); - move_tile_window( - window_i, - {number{}, - number<0>{}}); - return window_i; - }, - number{}); statically_indexed_array< - statically_indexed_array, + statically_indexed_array, NIterPerWarp> - b_warp_tensor_ping, b_warp_tensor_pong; + b_flat_dram_windows; + statically_indexed_array, + NIterPerWarp> + b_warp_tensor_ping; + statically_indexed_array, + NIterPerWarp> + b_warp_tensor_pong; - // pingpong buffer for Scale A and Scale B - auto scale_a_dram_window = make_tile_window( - scale_a_window.get_bottom_tensor_view(), - make_tuple(number{}, number<64 / WG::kM>{}), - scale_a_window.get_window_origin(), - PipelinePolicy::template MakeMXFP4_ScaleA_FlatDramTileDistribution()); - - auto scale_b_dram_window = make_tile_window( - scale_b_window.get_bottom_tensor_view(), - make_tuple(number{}, number<64 / WG::kN>{}), - scale_b_window.get_window_origin(), - PipelinePolicy::template MakeMXFP4_ScaleB_DramTileDistribution()); - - // ping pong buffer for scale A statically_indexed_array< - statically_indexed_array, - MIterPerWarp / MXdlPack> - scale_a_dram_windows; - statically_indexed_array, - MIterPerWarp / MXdlPack> - scale_a_tile_tensor_ping; - statically_indexed_array, - MIterPerWarp / MXdlPack> - scale_a_tile_tensor_pong; - - // ping pong buffer for scale B + statically_indexed_array, + ScaleNPerWarp> + scale_b_flat_dram_windows; statically_indexed_array< - statically_indexed_array, - NIterPerWarp / NXdlPack> - scale_b_dram_windows; - statically_indexed_array, - NIterPerWarp / NXdlPack> - scale_b_tile_tensor_ping; - statically_indexed_array, - NIterPerWarp / NXdlPack> - scale_b_tile_tensor_pong; + statically_indexed_array, + ScaleNPerWarp> + scale_b_warp_tensor_ping; + statically_indexed_array< + statically_indexed_array, + ScaleNPerWarp> + scale_b_warp_tensor_pong; - auto async_load_tile_ = [](auto lds, auto dram) { - async_load_tile(lds, dram, number<-1>{}, true_type{}, false_type{}); + 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 = + [&]([[maybe_unused]] 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, + "buffer_load_lds don't support element func fot A before mfma"); + }; +#else + auto prefill_lds_a_stage1 = + [&]([[maybe_unused]] 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 = [&]([[maybe_unused]] 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 - async_load_tile_(a_store_lds_window_ping, a_dram_window); - move_tile_window(a_dram_window, {0, kKPerBlock}); + prefill_lds_a_stage1(a_copy_lds_window_ping, a_copy_dram_window, number{}); + + // move A window to next k + move_tile_window(a_copy_dram_window, {0, kKPerBlock}); // prefetch B static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - b_warp_tensor_ping(nIter)(kIter) = load_tile_with_offset( - b_flat_dram_windows(nIter), number{}); - }); - // move B window to next flat K - move_tile_window(b_flat_dram_windows(nIter), {0, KIterPerWarp * KFlatPerBlockPerIter}); - }); + static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) + { + auto scale_n_iter = nIter / number{}; + auto scale_k_iter = kIter / number{}; - // prefetch Scale A - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; - move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), - {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); + scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = + 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_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = + load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); + } + auto packed_n_idx = nIter / number{}; + auto packed_n_rank = nIter % number{}; - scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) = - load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); + b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + move_tile_window(b_flat_dram_windows(nIter)(kIter), + {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + + packed_n_rank, + kIter * KFlatPerBlockPerIter}); + + ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); + b_warp_tensor_ping(nIter)(kIter) = ub.u; }); }); - // move Scale A window to next K - move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); + // 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}); - // prefetch Scale B - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; - move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), - {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); + prefill_lds_a_stage2(a_copy_lds_window_ping); - scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) = - load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); - }); - }); - // move Scale B window to next K - move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); __builtin_amdgcn_sched_barrier(0); // Prefetch A1 - if constexpr(HasHotLoop || TailNum == TailNumber::Even) - { - async_load_tile_(a_store_lds_window_pong, a_dram_window); - move_tile_window(a_dram_window, {0, kKPerBlock}); - } - // initialize C - clear_tile(c_block_tile); + prefill_lds_a_stage1(a_copy_lds_window_pong, a_copy_dram_window, number{}); + // move A window to next k + move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - statically_indexed_array a_warp_tensor; + // 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 - s_waitcnt_barrier(); - static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MXdlPack; - constexpr auto kIter = loadIter / MXdlPack; + statically_indexed_array{})(number<0>{}))), + m_preload> + a_warp_tensor; - a_warp_tensor(loadIter) = load_tile_with_offset( - a_warp_window_ping, tuple, number>{}); + static_for<0, m_preload, 1>{}([&](auto loadIter) { + constexpr auto mIter = loadIter % MIterPerWarp; + constexpr auto kIter = loadIter / MIterPerWarp; + a_warp_tensor(loadIter) = + load_tile(a_warp_windows_ping(number{})(number{})); }); __builtin_amdgcn_sched_barrier(0); - // MAIN LOOP - auto main_body_implx2 = [&]() mutable { - // prefetch B(2i+1) - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - b_warp_tensor_pong(nIter)(kIter) = load_tile_with_offset( - b_flat_dram_windows(nIter), number{}); - if constexpr(kIter == KIterPerWarp - 1) - move_tile_window(b_flat_dram_windows(nIter), - {0, BlockGemmShape::flatKPerBlock}); - }); + statically_indexed_array dequant_B_n; + + auto dequant_mxfp4 = [&](const auto& quant_weight_tensor, + const auto& scale_tensor, + auto xdl_nIter, + auto xdl_kIter) { + auto quant_idx_k = xdl_kIter % number{}; + + auto scale_idx_n = xdl_nIter % number{}; + auto scale_idx_k = (xdl_kIter % number{}) / number{}; + auto scale_offset = scale_idx_n + scale_idx_k * number{}; + + auto scale = scale_tensor.get_thread_buffer()[scale_offset]; + + constexpr int ScalarCnt = WG::BWarpTensor::get_thread_buffer_size(); + constexpr int PackedCnt = ScalarCnt / MXFP4PackedSize; + constexpr int float_mantissa = 23; + + uint32_t uscale = uint32_t(scale.data) << float_mantissa; + + using ComputeV2Type = + std::conditional_t, fp16x2_t, bf16x2_t>; + +#if defined(__gfx950__) + auto pk_mxfp4x4_to_compute_v2 = [](auto pk_mxfp4x4, float fscale, auto byte_idx) { + if constexpr(std::is_same_v) + { + return __builtin_amdgcn_cvt_scalef32_pk_f16_fp4( + pk_mxfp4x4, fscale, int(byte_idx)); + } + else if constexpr(std::is_same_v) + { + return __builtin_amdgcn_cvt_scalef32_pk_bf16_fp4( + pk_mxfp4x4, fscale, int(byte_idx)); + } + else + { + static_assert(sizeof(pk_mxfp4x4) == 0, "unsupported compute type"); + } + }; + static_for<0, PackedCnt, 1>{}([&](auto i) { + dequant_B_n[xdl_nIter].get_thread_buffer().template set_as( + i, + pk_mxfp4x4_to_compute_v2( + quant_weight_tensor[quant_idx_k], bit_cast(uscale), i)); }); - - // prefetch Scale A and Scale B (2i+1) - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; - move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), - {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); - - scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) = - load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); - }); +#else + auto pk_mxfp4_to_compute_v2 = [](auto pk_mxfp4, float fscale) { + if constexpr(std::is_same_v) + { + return pk_fp4_to_fp16x2(pk_mxfp4, fscale); + } + else if constexpr(std::is_same_v) + { + return pk_fp4_to_bf16x2(pk_mxfp4, fscale); + } + else + { + static_assert(sizeof(pk_mxfp4) == 0, "unsupported compute type"); + } + }; + static_for<0, PackedCnt, 1>{}([&](auto i) { + dequant_B_n[xdl_nIter].get_thread_buffer().template set_as( + i, + pk_mxfp4_to_compute_v2( + bit_cast>(quant_weight_tensor[quant_idx_k]) + .at(i), + bit_cast(uscale))); }); - - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; - move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), - {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); - - scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) = - load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); - }); - }); - - // GEMM 2i - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { - static_for<0, MXdlPack, 1>{}([&](auto imxdl) { - constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; - constexpr auto m_iter = mIter_pack * MXdlPack + imxdl; - constexpr auto k_iter = kIter_pack * KXdlPack + ikxdl; - static_for<0, NXdlPack, 1>{}([&](auto inxdl) { - constexpr auto n_iter = nIter_pack * NXdlPack + inxdl; - - // 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{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - - // warp GEMM - WG{}.template - operator()( - c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter_pack * number{} + inxdl)( - kIter_pack * number{} + ikxdl), - scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) - .get_thread_buffer()[0], - scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) - .get_thread_buffer()[0]); - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - constexpr auto addr = - m_iter % 2 + k_iter * 2 + m_iter / 2 * 4 + m_preload; - if constexpr(addr < (KIterPerWarp * MIterPerWarp) && - (nIter_pack == NIterPerWarp / NXdlPack - 1)) - { - constexpr auto AmIter = addr % 2 + addr / 4 * 2; - constexpr auto AkIter = addr / 2 % 2; - a_warp_tensor(number{}) = load_tile_with_offset( - a_warp_window_ping, - tuple, number>{}); - } - }); - }); - }); - }); - }); - // barrier as ds_load A(2i) and buffer_load_lds A(2i + 1) finished - s_waitcnt< // vmcnt - Bload_num + ScaleAload_num + ScaleBload_num>(); - block_sync_lds(); - - // Prefetch A(2i+2) - async_load_tile_(a_store_lds_window_ping, a_dram_window); - move_tile_window(a_dram_window, {0, kKPerBlock}); - - // move B window to next flat K - move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); - move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); - - // preload A(2i+1) - static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MXdlPack; - constexpr auto kIter = loadIter / MXdlPack; - a_warp_tensor(loadIter) = load_tile_with_offset( - a_warp_window_pong, tuple, number>{}); - }); - HotLoopScheduler(); - - ////////////////////////////// Next K ////////////////////////////// - - // prefetch B(2i+2) - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - b_warp_tensor_ping(nIter)(kIter) = load_tile_with_offset( - b_flat_dram_windows(nIter), number{}); - if constexpr(kIter == KIterPerWarp - 1) - move_tile_window(b_flat_dram_windows(nIter), - {0, BlockGemmShape::flatKPerBlock}); - }); - }); - - // prefetch Scale A and Scale B (2i+2) - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; - move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), - {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); - - scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) = - load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); - }); - }); - - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; - move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), - {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); - - scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) = - load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); - }); - }); - - // GEMM 2i+1 - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { - static_for<0, MXdlPack, 1>{}([&](auto imxdl) { - constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; - static_for<0, NXdlPack, 1>{}([&](auto inxdl) { - // 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{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - - // warp GEMM - WG{}.template - operator()( - c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_pong(nIter_pack * number{} + inxdl)( - kIter_pack * number{} + ikxdl), - scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) - .get_thread_buffer()[0], // scale A - scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) - .get_thread_buffer()[0]); // scale B - - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + - (kIter_pack * KXdlPack + ikxdl) * 2 + - (mIter_pack * MXdlPack + imxdl) / 2 * 4 + - m_preload; - if constexpr(addr < (KIterPerWarp * MIterPerWarp) && - (nIter_pack == NIterPerWarp / NXdlPack - 1)) - { - constexpr auto AmIter = addr % 2 + addr / 4 * 2; - constexpr auto AkIter = addr / 2 % 2; - a_warp_tensor(number{}) = load_tile_with_offset( - a_warp_window_pong, - tuple, number>{}); - } - }); - }); - }); - }); - }); - // barrier as ds_load A(2i + 1) and buffer_load_lds A(2i + 2) finished - s_waitcnt< // vmcnt - Bload_num + ScaleAload_num + ScaleBload_num>(); - block_sync_lds(); - - // Prefetch A(2i+3) - async_load_tile_(a_store_lds_window_pong, a_dram_window); - move_tile_window(a_dram_window, {0, kKPerBlock}); - // move B window to next flat K - move_tile_window(scale_a_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); - move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)}); - - // preload A(2i+2) - static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MXdlPack; - constexpr auto kIter = loadIter / MXdlPack; - a_warp_tensor(loadIter) = load_tile_with_offset( - a_warp_window_ping, tuple, number>{}); - }); - HotLoopScheduler(); +#endif }; - if constexpr(HasHotLoop) + // MAIN LOOP + index_t iCounter = (num_loop - 1) / 2; + while(iCounter > 0) { - index_t iCounter = (num_loop - 1) / 2; - do - { - main_body_implx2(); - iCounter--; - } while(iCounter > 0); + // prefetch B(2i+1) + static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) + { + auto scale_n_iter = nIter / number{}; + auto scale_k_iter = kIter / number{}; + + scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = + 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_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = + load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); + } + + auto packed_n_idx = nIter / number{}; + auto packed_n_rank = nIter % number{}; + + b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + + move_tile_window( + b_flat_dram_windows(nIter)(kIter), + {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + + packed_n_rank, + kIter * KFlatPerBlockPerIter}); + + ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); + b_warp_tensor_pong(nIter)(kIter) = ub.u; + }); + }); + + // Prefill A(2i+1) + prefill_lds_a_stage2(a_copy_lds_window_pong); + + // Prefetch A(2i+2) + prefill_lds_a_stage1( + a_copy_lds_window_ping, a_copy_dram_window, number{}); + // GEMM 2i + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + if constexpr(mIter == 0) + dequant_mxfp4( + b_warp_tensor_ping(nIter)(kIter / number{}), + scale_b_warp_tensor_ping(nIter / number{})( + kIter / number{}), + nIter, + kIter); + + // warp GEMM + WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + if constexpr((kIter * MIterPerWarp + mIter) < + (KIterPerWarp * MIterPerWarp - m_preload)) + { + constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; + constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); + a_warp_tensor(number{}) = + load_tile(a_warp_windows_ping(number{})(number{})); + } + + // 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{}); + + // 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}); + + static_for<0, m_preload, 1>{}([&](auto loadIter) { + constexpr auto mIter = loadIter % MIterPerWarp; + constexpr auto kIter = loadIter / MIterPerWarp; + a_warp_tensor(loadIter) = + load_tile(a_warp_windows_pong(number{})(number{})); + }); + HotLoopScheduler(); + + // Next K + + // prefetch B(2i+2) + static_for<0, MXFP4KPerWarp, 1>{}([&](auto kIter) { + static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) + { + auto scale_n_iter = nIter / number{}; + auto scale_k_iter = kIter / number{}; + + scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = + 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_b_warp_tensor_ping(scale_n_iter)(scale_k_iter) = + load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); + } + + auto packed_n_idx = nIter / number{}; + auto packed_n_rank = nIter % number{}; + + b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + move_tile_window( + b_flat_dram_windows(nIter)(kIter), + {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + + 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) + prefill_lds_a_stage2(a_copy_lds_window_ping); + + // Prefetch A(2i+3) + prefill_lds_a_stage1( + a_copy_lds_window_pong, a_copy_dram_window, number{}); + + // GEMM 2i+1 + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + if constexpr(mIter == 0) + dequant_mxfp4( + b_warp_tensor_pong(nIter)(kIter / number{}), + scale_b_warp_tensor_pong(nIter / number{})( + kIter / number{}), + nIter, + kIter); + + // warp GEMM + WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); + }); + // preload next A from lds + if constexpr((kIter * MIterPerWarp + mIter) < + (KIterPerWarp * MIterPerWarp - m_preload)) + { + constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; + constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); + a_warp_tensor(number{}) = + load_tile(a_warp_windows_pong(number{})(number{})); + } + + // 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{}); + + // 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}); + + static_for<0, m_preload, 1>{}([&](auto loadIter) { + constexpr auto mIter = loadIter % MIterPerWarp; + constexpr auto kIter = loadIter / MIterPerWarp; + a_warp_tensor(loadIter) = + load_tile(a_warp_windows_ping(number{})(number{})); + }); + HotLoopScheduler(); + + iCounter--; } // TAIL if constexpr(TailNum == TailNumber::Even) { // prefetch B(loopK) - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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) { - b_warp_tensor_pong(nIter)(kIter) = load_tile_with_offset( - b_flat_dram_windows(nIter), - make_tuple(number<0>{}, number{})); + if constexpr(nIter % XDL_PerScaleN == 0 && kIter % MXFP4K_PerScaleK == 0) + { + auto scale_n_iter = nIter / number{}; + auto scale_k_iter = kIter / number{}; + + scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter) = + 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_b_warp_tensor_pong(scale_n_iter)(scale_k_iter) = + load_tile(scale_b_flat_dram_windows(scale_n_iter)(scale_k_iter)); + } }); + + const int b_k_off_inter = kIter * kKPerBlock / MXFP4KPerWarp + b_k_off; + if(b_k_off_inter < kKPerBlock - k_padded_zeros) + { + static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + auto packed_n_idx = nIter / number{}; + auto packed_n_rank = nIter % number{}; + + b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + + move_tile_window( + b_flat_dram_windows(nIter)(kIter), + {packed_n_idx * ContinuousScaleNPerThread * NFlatPerBlockPerIter + + packed_n_rank, + kIter * KFlatPerBlockPerIter}); + + ub.mxfp4 = load_tile(b_flat_dram_windows(nIter)(kIter)); + b_warp_tensor_pong(nIter)(kIter) = ub.u; + }); + } }); - // prefetch Scale A and Scale B (2i+1) - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_a_dram_windows(mIter_pack)(kIter_pack) = scale_a_dram_window; - move_tile_window(scale_a_dram_windows(mIter_pack)(kIter_pack), - {mIter_pack * MWarp * WG::kM, kIter_pack * (64 / WG::kM)}); - - scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) = - load_tile(scale_a_dram_windows(mIter_pack)(kIter_pack)); - }); - }); - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - scale_b_dram_windows(nIter_pack)(kIter_pack) = scale_b_dram_window; - move_tile_window(scale_b_dram_windows(nIter_pack)(kIter_pack), - {nIter_pack * NWarp * WG::kN, kIter_pack * (64 / WG::kN)}); - - scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) = - load_tile(scale_b_dram_windows(nIter_pack)(kIter_pack)); - }); - }); + // Prefill A(loopK) + prefill_lds_a_stage2(a_copy_lds_window_pong); // GEMM loopK-1 - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { - static_for<0, MXdlPack, 1>{}([&](auto imxdl) { - constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; - static_for<0, NXdlPack, 1>{}([&](auto inxdl) { - // 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{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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; - // warp GEMM - WG{}.template - operator()( - c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter_pack * number{} + inxdl)( - kIter_pack * number{} + ikxdl), - scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) - .get_thread_buffer()[0], // scale A - scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) - .get_thread_buffer()[0]); // scale B + c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + - (kIter_pack * KXdlPack + ikxdl) * 2 + - (mIter_pack * MXdlPack + imxdl) / 2 * 4 + - m_preload; - if constexpr(addr < (KIterPerWarp * MIterPerWarp) && - (nIter_pack == NIterPerWarp / NXdlPack - 1)) - { - constexpr auto AmIter = addr % 2 + addr / 4 * 2; - constexpr auto AkIter = addr / 2 % 2; - a_warp_tensor(number{}) = load_tile_with_offset( - a_warp_window_ping, - tuple, number>{}); - } - }); - }); + if constexpr(mIter == 0) + dequant_mxfp4( + b_warp_tensor_ping(nIter)(kIter / number{}), + scale_b_warp_tensor_ping(nIter / number{})( + kIter / number{}), + nIter, + kIter); + + // warp GEMM + WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); }); + // preload next A from lds + if constexpr((kIter * MIterPerWarp + mIter) < + (KIterPerWarp * MIterPerWarp - m_preload)) + { + constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; + constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); + a_warp_tensor(number{}) = + load_tile(a_warp_windows_ping(number{})(number{})); + } + + // barrier + if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) + { + __builtin_amdgcn_s_waitcnt(Bload_total_num); + block_sync_lds(); + } }); }); - // barrier as ds_load A(2i) and buffer_load_lds A(2i + 1) finished - s_waitcnt< // vmcnt - Bload_num + ScaleAload_num + ScaleBload_num>(); - block_sync_lds(); - // preload A(2i+1) static_for<0, m_preload, 1>{}([&](auto loadIter) { - constexpr auto mIter = loadIter % MXdlPack; - constexpr auto kIter = loadIter / MXdlPack; - a_warp_tensor(loadIter) = load_tile_with_offset( - a_warp_window_pong, tuple, number>{}); + constexpr auto mIter = loadIter % MIterPerWarp; + constexpr auto kIter = loadIter / MIterPerWarp; + a_warp_tensor(loadIter) = + load_tile(a_warp_windows_pong(number{})(number{})); }); - Last2ndHotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + // Last2ndHotLoopScheduler(); // GEMM loopK - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { - static_for<0, MXdlPack, 1>{}([&](auto imxdl) { - constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; - static_for<0, NXdlPack, 1>{}([&](auto inxdl) { - // 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{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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; - // warp GEMM - WG{}.template - operator()( - c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_pong(nIter_pack * number{} + inxdl)( - kIter_pack * number{} + ikxdl), - scale_a_tile_tensor_pong(mIter_pack)(kIter_pack) - .get_thread_buffer()[0], // scale A - scale_b_tile_tensor_pong(nIter_pack)(kIter_pack) - .get_thread_buffer()[0]); // scale B + c_warp_tensor.get_thread_buffer() = + c_block_tile.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + - (kIter_pack * KXdlPack + ikxdl) * 2 + - (mIter_pack * MXdlPack + imxdl) / 2 * 4 + - m_preload; - if constexpr(addr < (KIterPerWarp * MIterPerWarp) && - (nIter_pack == NIterPerWarp / NXdlPack - 1)) - { - constexpr auto AmIter = addr % 2 + addr / 4 * 2; - constexpr auto AkIter = addr / 2 % 2; - a_warp_tensor(number{}) = load_tile_with_offset( - a_warp_window_pong, - tuple, number>{}); - } - }); + if constexpr(mIter == 0) + dequant_mxfp4( + b_warp_tensor_pong(nIter)(kIter / number{}), + scale_b_warp_tensor_pong(nIter / number{})( + kIter / number{}), + nIter, + kIter); + + // warp GEMM + WG{}(c_warp_tensor, + a_warp_tensor(number{}), + dequant_B_n[nIter]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); }); + if constexpr((kIter * MIterPerWarp + mIter) < + (KIterPerWarp * MIterPerWarp - m_preload)) + { + constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; + constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); + a_warp_tensor(number{}) = + load_tile(a_warp_windows_pong(number{})(number{})); + } + // barrier + // if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) + // { + // block_sync_lds(); + // } }); - }); + } }); LastHotLoopScheduler(); } else if constexpr(TailNum == TailNumber::Odd) { // GEMM loopK - static_for<0, KIterPerWarp / KXdlPack, 1>{}([&](auto kIter_pack) { - static_for<0, MIterPerWarp / MXdlPack, 1>{}([&](auto mIter_pack) { - static_for<0, NIterPerWarp / NXdlPack, 1>{}([&](auto nIter_pack) { - static_for<0, KXdlPack, 1>{}([&](auto ikxdl) { - static_for<0, MXdlPack, 1>{}([&](auto imxdl) { - constexpr auto AwarpIter = imxdl + ikxdl * MXdlPack; - static_for<0, NXdlPack, 1>{}([&](auto inxdl) { - // 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{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + 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; - // warp GEMM - WG{}.template - operator()( - c_warp_tensor, - a_warp_tensor(number{}), - b_warp_tensor_ping(nIter_pack * number{} + inxdl)( - kIter_pack * number{} + ikxdl), - scale_a_tile_tensor_ping(mIter_pack)(kIter_pack) - .get_thread_buffer()[0], // scale A - scale_b_tile_tensor_ping(nIter_pack)(kIter_pack) - .get_thread_buffer()[0]); // scale B + c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - // write C warp tensor into C block tensor - c_block_tile.set_y_sliced_thread_data( - merge_sequences(sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), - c_warp_tensor.get_thread_buffer()); - }); - // preload next A from lds - constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 + - (kIter_pack * KXdlPack + ikxdl) * 2 + - (mIter_pack * MXdlPack + imxdl) / 2 * 4 + - m_preload; - if constexpr(addr < (KIterPerWarp * MIterPerWarp) && - (nIter_pack == NIterPerWarp / NXdlPack - 1)) - { - constexpr auto AmIter = addr % 2 + addr / 4 * 2; - constexpr auto AkIter = addr / 2 % 2; - a_warp_tensor(number{}) = load_tile_with_offset( - a_warp_window_ping, - tuple, number>{}); - } - }); - }); + if constexpr(mIter == 0) + dequant_mxfp4( + b_warp_tensor_ping(nIter)(kIter / number{}), + scale_b_warp_tensor_ping(nIter / number{})( + kIter / number{}), + nIter, + kIter); + // warp GEMM + WG{}(c_warp_tensor, a_warp_tensor(number{}), dequant_B_n[nIter]); + + // write C warp tensor into C block tensor + c_block_tile.set_y_sliced_thread_data( + merge_sequences(sequence{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + c_warp_tensor.get_thread_buffer()); }); + // preload next A from lds + if constexpr((kIter * MIterPerWarp + mIter) < + (KIterPerWarp * MIterPerWarp - m_preload)) + { + constexpr auto AmIter = (mIter + m_preload) % MIterPerWarp; + constexpr auto AkIter = (kIter + (mIter + m_preload) / MIterPerWarp); + a_warp_tensor(number{}) = + load_tile(a_warp_windows_ping(number{})(number{})); + } + + // barrier + if constexpr((kIter == KIterPerWarp - 1) && (mIter == MIter_2nd_last)) + { + __builtin_amdgcn_s_waitcnt(Bload_total_num); + block_sync_lds(); + } }); }); LastHotLoopScheduler(); } - else - { - static_assert(false, "Wrong TailNum"); - } + return c_block_tile; } + + template + 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, + const index_t k_padded_zeros, + 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, + k_padded_zeros, + p_smem_ping, + p_smem_pong); + } + + template + 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); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index cf32ea069c..217f5ff692 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -248,7 +248,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy template CK_TILE_HOST_DEVICE static constexpr auto - TransformF16xF4_ATensorView(const NativeADramTensorView& a_dram_view) + TransformF8_ATensorView(const NativeADramTensorView& a_dram_view) { #if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS constexpr int DynamicTileOffsetFlag = 0; @@ -262,7 +262,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; constexpr index_t KPack = GetSmemPackA(); - constexpr int ContiguousThreadsCntInDS_READ_16B = 4; + constexpr int ContiguousThreadsCntInDS_READ_16B = 2; // implement swizzle pattern on global side // because we can't adjust the ds_write pattern of BUFFER_LOAD_LDS. @@ -305,7 +305,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy } template - CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ReadALdsBlockDescriptor() + CK_TILE_HOST_DEVICE static constexpr auto MakeF8_ReadALdsBlockDescriptor() { constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); @@ -323,7 +323,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy number{}, number<1>{}); - constexpr int ContiguousThreadsCntInDS_READ_16B = 4; + constexpr int ContiguousThreadsCntInDS_READ_16B = 2; constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor( a_lds_block_desc_0, @@ -345,7 +345,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy } template - CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_WriteALdsBlockDescriptor() + CK_TILE_HOST_DEVICE static constexpr auto MakeF8_WriteALdsBlockDescriptor() { #if CKTILE_FLATMM_USE_BUFFER_LOAD_LDS constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; @@ -361,7 +361,7 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy } template - CK_TILE_HOST_DEVICE static constexpr auto MakeF16xF4_ALDS_TileDistribution() + CK_TILE_HOST_DEVICE static constexpr auto MakeF8_ALDS_TileDistribution() { using TileShape = typename Problem::BlockGemmShape; @@ -385,72 +385,6 @@ struct F8xMXF4MoeFlatmmPipelineAgBgCrPolicy : F16xMXF4FlatmmPipelineAgBgCrPolicy sequence<2>, sequence<2>>{}); } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeFp4BFlatDramTileDistribution() - { - using TileShape = typename Problem::BlockGemmShape; - - static_assert(TileShape::WarpTile::at(I1) == 16, "only for XDL_N == 16"); - - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t WaveSize = get_warp_size(); - constexpr index_t WaveNum = BlockSize / WaveSize; - - constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim - constexpr index_t KWavePerBlk = 1; - - constexpr index_t NWavePerBlk = TileShape::BlockWarps::at(number<1>{}); // N_Warp - - constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp; - - return make_static_tile_distribution( - tile_distribution_encoding< - sequence, // ? - tuple, // second - // direction - sequence>, // first direction - // wave in blk, // thd in wave - // // - tuple, sequence<2>>, // which direction - tuple, sequence<1>>, // which index - // - sequence<2>, - sequence<2>>{}); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeFp4ScaleBFlatDramTileDistribution() - { - using TileShape = typename Problem::BlockGemmShape; // ck_tile::TileFlatmmShape - - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t WaveSize = get_warp_size(); - [[maybe_unused]] constexpr index_t WaveNum = BlockSize / WaveSize; - - constexpr index_t N_Warp = TileShape::BlockWarps::at(number<1>{}); - - [[maybe_unused]] constexpr index_t XDLPerBlock = - TileShape::kK / TileShape::WarpTile::at(I2); - constexpr index_t K_Lane = 64 / TileShape::WarpTile::at(I1); - constexpr index_t N_Lane = TileShape::WarpTile::at(I1); - - constexpr index_t NWavePerBlk = N_Warp; - - return make_static_tile_distribution( - tile_distribution_encoding< - sequence<>, // ? - tuple, // second direction - sequence>, // first - // direction - // wave in blk, // thd in wave - // // - tuple, sequence<2, 2>>, // which direction - tuple, sequence<0, 1>>, // which index - // - sequence<2>, - sequence<2>>{}); - } }; } // namespace ck_tile