diff --git a/example/67_gemm_microscaling/CMakeLists.txt b/example/67_gemm_microscaling/CMakeLists.txt index 6c0c973880..9731fe7376 100644 --- a/example/67_gemm_microscaling/CMakeLists.txt +++ b/example/67_gemm_microscaling/CMakeLists.txt @@ -15,4 +15,8 @@ add_example_dependencies(example_gemm_mx example_gemm_mx_fp4) set(FP4_MXGEMM_OPTIONS) list(APPEND FP4_MXGEMM_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1") list(APPEND FP4_MXGEMM_OPTIONS -v --save-temps -Wno-gnu-line-marker -ftemplate-backtrace-limit=0) -target_compile_options(example_gemm_mx_fp4 PRIVATE ${FP4_MXGEMM_OPTIONS}) \ No newline at end of file +set(FP8_MXGEMM_OPTIONS) +list(APPEND FP8_MXGEMM_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1 -mllvm --slp-threshold=-32") +list(APPEND FP8_MXGEMM_OPTIONS -v --save-temps -Wno-gnu-line-marker -ftemplate-backtrace-limit=0) +target_compile_options(example_gemm_mx_fp4 PRIVATE ${FP4_MXGEMM_OPTIONS}) +target_compile_options(example_gemm_mx_fp8 PRIVATE ${FP8_MXGEMM_OPTIONS}) \ No newline at end of file diff --git a/example/67_gemm_microscaling/gemm_mx_fp4.cpp b/example/67_gemm_microscaling/gemm_mx_fp4.cpp index b913671ab9..84c12bd0ac 100644 --- a/example/67_gemm_microscaling/gemm_mx_fp4.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp4.cpp @@ -23,19 +23,16 @@ using AElementOp = PassThrough; // elementwise transformation for A matrix using BElementOp = PassThrough; // elementwise transformation for B matrix using CElementOp = PassThrough; // elementwise transformation for C matrix +constexpr ck::index_t DataPackedSize = 2; // Packed representation of data constexpr ck::index_t ScaleBlockSize = 32; // scaling block size -constexpr ck::index_t KPerBlock = 256; +constexpr ck::index_t KPerBlock = 256 / DataPackedSize; // 256 f4 = 128 fp4x2 constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; constexpr auto BlkGemmPSched = ck::BlockGemmPipelineScheduler::Intrawave; constexpr auto BlkGemmPVer = ck::BlockGemmPipelineVersion::v3; -// v3 should be performant one, However -// 1. some bug existed cause memory access fault in some cases, MNK=2k2k2k -// 2. Register spill observed, most likely unpack the e8m0 from single register then feed to -// scaled mfma. // AB DataType: f4x2_pk_t -// Mathmatically, all numbers are represented as f4. +// Mathmatically, all numbers are represented as f4x2. using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffleV3< ALayout, // ALayout BLayout, // BLayout @@ -56,8 +53,8 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffle 128, // MPerBlock 256, // NPerBlock KPerBlock, // KPerBlock - 32, // AK1 - 32, // BK1 + 16, // AK1 + 16, // BK1 16, // MPerXDL 16, // NPerXDL 4, // MXdlPerWave @@ -66,15 +63,15 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffle S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder S<1, 0, 2>, // ABlockTransferSrcAccessOrder 2, // ABlockTransferSrcVectorDim - 32, // ABlockTransferSrcScalarPerVector - 32, // ABlockTransferDstScalarPerVector_AK1 + 16, // ABlockTransferSrcScalarPerVector + 16, // ABlockTransferDstScalarPerVector_AK1 false, // ABlockLdsExtraM S<8, 32, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 S<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder S<1, 0, 2>, // BBlockTransferSrcAccessOrder 2, // BBlockTransferSrcVectorDim - 32, // BBlockTransferSrcScalarPerVector - 32, // BBlockTransferDstScalarPerVector_BK1 + 16, // BBlockTransferSrcScalarPerVector + 16, // BBlockTransferDstScalarPerVector_BK1 false, // BBlockLdsExtraN 2, // CShuffleMXdlPerWavePerShuffle 2, // CShuffleNXdlPerWavePerShuffle diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp index dd74c7e51c..9afdb29892 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp @@ -35,6 +35,11 @@ struct BlockwiseGemmXdlops_mx_pipeline_base using ComputeTypeB = BDataType; using AccType = float; // for now only support V_MFMA_SCALE_F32 + static constexpr index_t APackedSize = + is_same_v, f4x2_pk_t> ? 2 : 1; + static constexpr index_t BPackedSize = + is_same_v, f4x2_pk_t> ? 2 : 1; + static constexpr auto I0 = Number<0>{}; static constexpr auto I1 = Number<1>{}; static constexpr auto I2 = Number<2>{}; @@ -51,19 +56,14 @@ struct BlockwiseGemmXdlops_mx_pipeline_base static constexpr index_t B_K1 = BTileDesc{}.GetLength(I2); static constexpr auto xdlops_gemm = - XdlopsGemm{}; + XdlopsGemm{}; static constexpr index_t AMmaKStride = KPack; static constexpr index_t BMmaKStride = KPack; //> store rows/cols into thread registers in chunks of 16 //> e.g. [k0,...,k15,k64,...,k79] or [k0,...,k15,k32,...,k47] - static constexpr index_t APackedSize = - is_same_v, f4x2_pk_t> ? 2 : 1; - static constexpr index_t BPackedSize = - is_same_v, f4x2_pk_t> ? 2 : 1; - - static constexpr index_t KThreadChunk = 16 * APackedSize / sizeof(ComputeTypeA); + static constexpr index_t KThreadChunk = 16 / sizeof(ComputeTypeA); static constexpr index_t KPerThread = KPerBlock / xdlops_gemm.K0PerXdlops; static constexpr index_t KRepeat = KPerThread / KPack; @@ -381,11 +381,11 @@ struct BlockwiseGemmXdlops_mx_pipeline_base I1, Number{}, Number{}, - Number{}), - make_tuple(Number{}, - Number{}, - Number{}, - Number{}, + Number{}), + make_tuple(Number{}, + Number{}, + Number{}, + Number{}, I1)); // B[N0, N1, N2, KPack] @@ -394,11 +394,11 @@ struct BlockwiseGemmXdlops_mx_pipeline_base I1, Number{}, Number{}, - Number{}), - make_tuple(Number{}, - Number{}, - Number{}, - Number{}, + Number{}), + make_tuple(Number{}, + Number{}, + Number{}, + Number{}, I1)); // C[M, N, NumRegXdlops] diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp index 91b6aa2a5e..9f9a104fe5 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp @@ -163,7 +163,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx How many mx-vectors in each row/col is processed in one call to xdlops_gemm.Run() - static constexpr auto ScalesPerXdlopsRun = (KPack * xdlops_gemm.K0PerXdlops) / ScaleBlockSize; + static constexpr auto ScalesPerXdlopsRun = (APackedSize * KPack * xdlops_gemm.K0PerXdlops) / ScaleBlockSize; //> How many scales a thread must read to accommodate one call to xdlops_gemm.Run() static constexpr auto ScalesPerXdlopsRunPerThread = @@ -194,11 +194,11 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx - // sizeof(ComputeDataType) / sizeof(BDataType) - // ? sizeof(ComputeDataType) / sizeof(ADataType) - // : sizeof(ComputeDataType) / sizeof(BDataType); constexpr auto num_mfma_stage1 = num_mfma_inst - (num_dsread_a_mfma + num_dsread_b_mfma); constexpr auto num_mfma_per_issue = num_mfma_stage1 / (num_buffer_load_inst_a + num_buffer_load_inst_b); @@ -430,9 +425,9 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto k) { - constexpr auto k_step = k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); + constexpr auto k_step = k * xdlops_gemm.KPerXdlops/APackedSize * (APackedSize * KPack / xdlops_gemm.K1PerXdlops); static_for<0, MRepeat, 1>{}([&](auto m0) { - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / (APackedSize * KThreadChunk), 1>{}([&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; a_thread_copy_.Run(a_block_desc_m0_m1_m2_m3_k, @@ -453,7 +448,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto n0) { // read block data in chunks to assemble correct thread vectors - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / (BPackedSize * KThreadChunk), 1>{}([&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; b_thread_copy_.Run(b_block_desc_n0_n1_n2_n3_k, @@ -574,12 +569,12 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - vector_type + vector_type a_thread_vec; - vector_type + vector_type b_thread_vec; - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + static_for<0, KPack, 1>{}([&](auto ik) { a_thread_vec.template AsType()( ik) = a_thread_buf [Number{}([&](auto k) { constexpr auto k_step = - k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); + k * xdlops_gemm.KPerXdlops/APackedSize * (APackedSize * KPack / xdlops_gemm.K1PerXdlops); static_for<0, MRepeat, 1>{}([&](auto m0) { - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}( + static_for<0, xdlops_gemm.K1PerXdlops / (APackedSize * KThreadChunk), 1>{}( [&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * @@ -668,7 +663,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto n0) { // read block data in chunks to assemble correct thread vectors - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}( + static_for<0, xdlops_gemm.K1PerXdlops / (BPackedSize * KThreadChunk), 1>{}( [&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * @@ -772,10 +767,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - vector_type a_thread_vec; - vector_type b_thread_vec; + vector_type a_thread_vec; + vector_type b_thread_vec; - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + static_for<0, KPack, 1>{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; @@ -825,9 +820,9 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto k) { constexpr auto k_step = - k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); + k * xdlops_gemm.KPerXdlops/APackedSize * (APackedSize * KPack / xdlops_gemm.K1PerXdlops); static_for<0, MRepeat, 1>{}([&](auto m0) { - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / (APackedSize * KThreadChunk), 1>{}([&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; a_thread_copy_.Run(a_block_desc_m0_m1_m2_m3_k, @@ -848,7 +843,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto n0) { // read block data in chunks to assemble correct thread vectors - static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { + static_for<0, xdlops_gemm.K1PerXdlops / (BPackedSize * KThreadChunk), 1>{}([&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; b_thread_copy_.Run(b_block_desc_n0_n1_n2_n3_k, @@ -900,10 +895,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - vector_type a_thread_vec; - vector_type b_thread_vec; + vector_type a_thread_vec; + vector_type b_thread_vec; - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + static_for<0, KPack, 1>{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; @@ -982,16 +977,20 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - vector_type a_thread_vec; - vector_type b_thread_vec; + vector_type a_thread_vec; + vector_type b_thread_vec; - static_for<0, KPack / APackedSize, 1>{}([&](auto ik) { + static_for<0, KPack, 1>{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; b_thread_vec.template AsType()(ik) = b_thread_buf[Number{}]; + CK_PRINT, + Number, + Number + >(); }); using mfma_input_type_a = diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp index df9ed8b334..699c6c40cb 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp @@ -220,6 +220,14 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX, pk_i4_t> || + is_same_v, f4x2_pk_t>) + return 2; + else + return 1; + }(); + if(stream_config.log_level_ > 0) { arg.Print(); @@ -296,18 +304,10 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX, pk_i4_t> || - is_same_v, f4x2_pk_t>) - return 2; - else - return 1; - }(); - constexpr index_t minimum_occupancy = BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave ? (BlkGemmPipelineVer == BlockGemmPipelineVersion::v3 && - MPerBlock * NPerBlock * KPerBlock * sizeof(ADataType) / APackedSize <= + MPerBlock * NPerBlock * KPerBlock * sizeof(ADataType) <= 128 * 128 * 64 * 2) ? 2 : 1 @@ -418,31 +418,6 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX; - Run(kernel); - } - else - { - const auto kernel = - kernel_gemm_xdl_cshuffle_v3; - Run(kernel); - } -#endif - } } return ave_time; diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp index 797008c1c1..8e9299ff52 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp @@ -174,16 +174,6 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // Should be a multiple of k_per_blk. // TODO: Move this to blockwise pipeline base // KPack in packed data types for pk A/B - static constexpr index_t KPack = - math::max(lcm_AK1_BK1, - MfmaSelector::selected_mfma.k_per_blk); - - using ThisThreadBlock = ThisThreadBlock; static constexpr index_t APackedSize = []() { if constexpr(is_same_v, pk_i4_t> || @@ -201,6 +191,17 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 return 1; }(); + static constexpr index_t KPack = + math::max(lcm_AK1_BK1, + MfmaSelector::selected_mfma.k_per_blk/APackedSize); + + using ThisThreadBlock = ThisThreadBlock; + __host__ static auto CalculateGridSize(index_t M, index_t N, index_t KBatch) { return std::make_tuple(Block2CTileMap::CalculateGridSize(M, N), 1, KBatch); @@ -647,10 +648,10 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 bool is_reduce_ = false) : Problem{M_, N_, - K_, - StrideA_, + K_/APackedSize, + StrideA_/APackedSize, StrideScaleA_, - StrideB_, + StrideB_/BPackedSize, StrideScaleB_, StrideC_, k_batch_}, @@ -695,7 +696,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 { if constexpr(is_same_v) { - a_k_split_offset = k_id * karg.KRead / APackedSize; + a_k_split_offset = k_id * karg.KRead; } else if constexpr(is_same_v) { @@ -710,33 +711,33 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 { if constexpr(!PermuteB) { - b_k_split_offset = k_id * karg.KRead / BPackedSize; + b_k_split_offset = k_id * karg.KRead; } else { const int k0_offset = karg.KRead * karg.N; - b_k_split_offset = k_id * k0_offset / BPackedSize; + b_k_split_offset = k_id * k0_offset; } } // Calculate A scale offset if constexpr(is_same_v) { - a_scale_k_split_offset = k_id * karg.KRead / ScaleBlockSize; + a_scale_k_split_offset = k_id * karg.KRead / (ScaleBlockSize/APackedSize); } else if constexpr(is_same_v) { - a_scale_k_split_offset = k_id * karg.KRead / ScaleBlockSize * karg.StrideScaleA; + a_scale_k_split_offset = k_id * karg.KRead / (ScaleBlockSize/APackedSize) * karg.StrideScaleA; } // Calculate B scale offset if constexpr(is_same_v) { - b_scale_k_split_offset = k_id * (karg.KRead / ScaleBlockSize) * karg.StrideScaleB; + b_scale_k_split_offset = k_id * (karg.KRead / (ScaleBlockSize/BPackedSize)) * karg.StrideScaleB; } else if constexpr(is_same_v) { - b_scale_k_split_offset = k_id * karg.KRead / ScaleBlockSize; + b_scale_k_split_offset = k_id * karg.KRead / (ScaleBlockSize/BPackedSize); } if(k_id < (karg.KBatch - 1)) @@ -1061,8 +1062,8 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 constexpr auto c_block_size = c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize(); - return math::max((a_block_space_size_aligned * sizeof(ADataType) / APackedSize + - b_block_space_size_aligned * sizeof(BDataType) / BPackedSize), + return math::max((a_block_space_size_aligned * sizeof(ADataType) + + b_block_space_size_aligned * sizeof(BDataType)), c_block_size * sizeof(CShuffleDataType)); } @@ -1073,7 +1074,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 (NPerBlock % (NXdlPerWave * NPerXdl)) == 0, "Invalid tuning param!"); - static_assert(KPerBlock % ScaleBlockSize == 0, + static_assert(KPerBlock % (ScaleBlockSize/BPackedSize) == 0, "KPerBlock should be multiple of ScaleBlockSize"); if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding || @@ -1449,13 +1450,12 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // Cast after lds auto a_block_buf = make_dynamic_buffer( static_cast(p_shared), - a_block_desc_ak0_m_ak1.GetElementSpaceSize() / APackedSize); + a_block_desc_ak0_m_ak1.GetElementSpaceSize()); auto b_block_buf = make_dynamic_buffer( reinterpret_cast(static_cast(p_shared) + a_block_space_size_aligned * - sizeof(ADataType) / - APackedSize), - b_block_desc_bk0_n_bk1.GetElementSpaceSize() / BPackedSize); + sizeof(ADataType)), + b_block_desc_bk0_n_bk1.GetElementSpaceSize()); constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0); constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock / BK1Number, 0, 0); @@ -1799,27 +1799,16 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( c_grid_desc_m_n, problem.MBlock, problem.NBlock); -#if 0 - // A Scale grid - const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor( - make_tuple(problem.M, math::integer_divide_ceil(problem.K, ScaleBlockSize)), - make_tuple(problem.StrideScaleA, 1)); - - // B Scale grid transposed - const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor( - make_tuple(problem.N, math::integer_divide_ceil(problem.K, ScaleBlockSize)), - make_tuple(problem.StrideScaleB, 1)); -#endif // A/B shuffled scale for better 8-bit scale access pattern // MNRepeat -> KRepeat -> KThreadPerXdl -> MNThreadPerXdl -> KXdlPack -> MNXdlPack const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor_packed(make_tuple( problem.M / (MXdlPack * MPerXdl), - math::integer_divide_ceil(problem.K, ScaleBlockSize) / (KXdlPack * 64 / MPerXdl), + math::integer_divide_ceil(problem.K, (ScaleBlockSize/APackedSize)) / (KXdlPack * 64 / MPerXdl), 64 * KXdlPack * MXdlPack / scale_pack_size_a)); const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor_packed(make_tuple( problem.N / (NXdlPack * NPerXdl), - math::integer_divide_ceil(problem.K, ScaleBlockSize) / (KXdlPack * 64 / NPerXdl), + math::integer_divide_ceil(problem.K, (ScaleBlockSize/BPackedSize)) / (KXdlPack * 64 / NPerXdl), 64 * KXdlPack * NXdlPack / scale_pack_size_b)); Run( bit_cast(static_cast(p_shared_0) + - a_block_space_size_aligned * sizeof(ADataType) / APackedSize), + a_block_space_size_aligned * sizeof(ADataType)), b_block_desc_bk0_n_bk1.GetElementSpaceSize()); auto a_block_buf_pong = make_dynamic_buffer( @@ -1999,7 +1988,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 auto b_block_buf_pong = make_dynamic_buffer( bit_cast(bit_cast(p_shared_1) + - a_block_space_size_aligned * sizeof(ADataType) / APackedSize), + a_block_space_size_aligned * sizeof(ADataType)), b_block_desc_bk0_n_bk1.GetElementSpaceSize()); auto a_block_bufs = make_tuple(a_block_buf_ping, a_block_buf_pong); @@ -2026,9 +2015,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 static constexpr auto KPerThread = KPerBlock / K0PerXdlops; const index_t ScaleSliceSizeN = NXdlPerWave; - static constexpr auto ScaleSliceSizeK = (KPerThread + ScaleBlockSize - 1) / ScaleBlockSize; + static constexpr auto ScaleSliceSizeK = (KPerThread + (ScaleBlockSize/BPackedSize) - 1) / (ScaleBlockSize/BPackedSize); static constexpr auto KBlockScaleSliceSizeK = - (KPerBlock + ScaleBlockSize - 1) / ScaleBlockSize; + (KPerBlock + (ScaleBlockSize/BPackedSize) - 1) / (ScaleBlockSize/BPackedSize); constexpr auto b_scale_thread_desc = make_naive_tensor_descriptor_packed( make_tuple(Number{}, Number{})); @@ -2054,7 +2043,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 false>( b_scale_grid_desc_bn_ak, make_multi_index(block_n_id * NPerBlock + b_thread_offset_n, - b_thread_offset_k / ScaleBlockSize)); + b_thread_offset_k / (ScaleBlockSize/BPackedSize))); constexpr auto b_scale_thread_slice_copy_step = make_tuple(make_multi_index(NWaves * NPerXdl, 0), @@ -2303,7 +2292,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 c_grid_desc_m_n, problem.MBlock, problem.NBlock); const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor( - make_tuple(problem.N, math::integer_divide_ceil(problem.K, ScaleBlockSize)), + make_tuple(problem.N, math::integer_divide_ceil(problem.K, ScaleBlockSize/BPackedSize)), make_tuple(problem.StrideScaleB, 1)); Run_2Lds, pk_i4_t> || - is_same_v, f4x2_pk_t>) + if constexpr(is_same_v, pk_i4_t>) return 2; else return 1; @@ -1043,8 +1042,7 @@ struct ThreadwiseTensorSliceTransfer_v4 using SrcCoordStep = decltype(make_tensor_coordinate_step(SrcDesc{}, Index{})); static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t> || - is_same_v, f4x2_pk_t>) + if constexpr(is_same_v, pk_i4_t>) return 2; else return 1; @@ -1175,9 +1173,6 @@ struct ThreadwiseTensorSliceTransfer_v4 src_tmp_vector.template AsType()(Number<0>{}) = src_buf.template Get(src_data_coord.GetOffset() / PackedSize, is_src_valid); - // printf("TID%03d GetOffset() / PackedSize = %d\n", - // get_thread_local_1d_id(), - // src_data_coord.GetOffset() / PackedSize); } else if constexpr(SrcBuffer::IsStaticBuffer()) { @@ -1510,8 +1505,7 @@ struct ThreadwiseTensorSliceTransfer_StaticToStatic using Index = MultiIndex; static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t> || - is_same_v, f4x2_pk_t>) + if constexpr(is_same_v, pk_i4_t>) return 2; else return 1; diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp index 9e988f61c0..a18e75dc95 100644 --- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp @@ -69,8 +69,7 @@ struct ThreadwiseTensorSliceTransfer_v3r1 static constexpr auto I16 = Number<16>{}; static constexpr index_t PackedSize = []() { - if constexpr(is_same_v, pk_i4_t> || - is_same_v, f4x2_pk_t>) + if constexpr(is_same_v, pk_i4_t>) return 2; else return 1;