diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_gather_direct_load.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_gather_direct_load.hpp index 8ad8e3b94d..0944151f4f 100644 --- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_gather_direct_load.hpp +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_gather_direct_load.hpp @@ -280,13 +280,25 @@ struct ThreadGroupTensorSliceTransfer_Gather_DirectLoad dst_offset, // *(reinterpret_cast(&(dst_buf[dst_offset + 0].data))), *(reinterpret_cast( - &(dst_buf[dst_offset + 16 * threadIdx.x].data))), + &(dst_buf[dst_offset + 0 + 16 * threadIdx.x].data))), *(reinterpret_cast( - &(dst_buf[dst_offset + 16 * threadIdx.x].data))), + &(dst_buf[dst_offset + 4 + 16 * threadIdx.x].data))), *(reinterpret_cast( - &(dst_buf[dst_offset + 32 * threadIdx.x].data))), + &(dst_buf[dst_offset + 8 + 16 * threadIdx.x].data))), *(reinterpret_cast( - &(dst_buf[dst_offset + 48 * threadIdx.x].data)))); + &(dst_buf[dst_offset + 12 + 16 * threadIdx.x].data)))); + +#else + __builtin_amdgcn_s_waitcnt(3952); + block_sync_lds(); + printf("blkx: %u, blky: %u, tid: %u, red_id: %d src: %d (cal: %d, gather: %d)\n", + blockIdx.x, + blockIdx.y, + threadIdx.x, + static_cast(ordered_dst_access_idx[Number{}]), + src_offset, + src_coord_.GetOffset(), + gather_offset); #endif constexpr auto move_src_on_dim = [&]() constexpr @@ -370,30 +382,6 @@ struct ThreadGroupTensorSliceTransfer_Gather_DirectLoad }); }); -#if 0 - __builtin_amdgcn_s_waitcnt(3952); - block_sync_lds(); - - if(threadIdx.x == 0) - { - // Print the contents of the destination buffer. - printf("blkx: %u, blky: %u, tid: %u, a_dst_buf_offset=<%d, %d, %d, %d>, " - "a_dst_buffer=<%02x, %02x, %02x, %02x>\n", - blockIdx.x, - blockIdx.y, - threadIdx.x, - 0, - 16, - 32, - 48, - static_cast(dst_buf[Number<0>{}].data), - static_cast(dst_buf[Number<16>{}].data), - static_cast(dst_buf[Number<32>{}].data), - static_cast(dst_buf[Number<48>{}].data)); - } - -#endif - // Reset the destination slice since the entire buffer has been already filled. ResetDstSliceWindow(dst_desc); } diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_mx_gemm_bns.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_mx_gemm_bns.hpp index c49996c4d5..dfaf185751 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_moe_mx_gemm_bns.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_mx_gemm_bns.hpp @@ -277,20 +277,20 @@ struct DeviceMoeGemmMXBNS : public DeviceMoEGemmMXBPreShuffle; + const auto kernel = kernel_moe_mxgemm_2lds; RunKernel(kernel); } else { - const auto kernel = kernel_moe_mxgemm; + const auto kernel = kernel_moe_mxgemm_2lds; RunKernel(kernel); } } @@ -327,20 +327,20 @@ struct DeviceMoeGemmMXBNS : public DeviceMoEGemmMXBPreShuffle; + const auto kernel = kernel_moe_mxgemm_2lds; RunKernel(kernel); } else { - const auto kernel = kernel_moe_mxgemm; + const auto kernel = kernel_moe_mxgemm_2lds; RunKernel(kernel); } } diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm_bns.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm_bns.hpp index 46b71737f0..57c6e0efe5 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm_bns.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm_bns.hpp @@ -35,6 +35,7 @@ enum Activation silu_and_mul = 1 }; +#if 0 template __global__ void #if CK_USE_LAUNCH_BOUNDS - __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) +__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) #endif // __attribute__((amdgpu_waves_per_eu(1, 1))) kernel_moe_mxgemm(typename GridwiseGemm::Argument karg) @@ -71,6 +72,7 @@ __global__ void ignore = karg; #endif // end of if (defined(__gfx9__)) } +#endif template @@ -1488,10 +1491,10 @@ struct GridwiseMoeGemmMXBNS { token_offset = token_offset * problem.TopK + (fused_token >> 24); } - gather_offsets(m0) = static_cast(token_offset) * problem.K; + gather_offsets(m0) = static_cast(token_offset); }); -#if 0 +#if 1 printf("blkx: %u, blky: %u, tidx: %u, token_pos: %d, gather_offsets:<%d, %d, %d, %d>\n", blockIdx.x, blockIdx.y, @@ -1967,18 +1970,16 @@ struct GridwiseMoeGemmMXBNS // tuple of reference to C/Ds tensor descriptors const auto c_ds_desc_refs = concat_tuple_of_reference( tie(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), - generate_tie( - [&](auto i) -> const auto& // return type should be reference - { return ds_grid_desc_mblock_mperblock_nblock_nperblock[i]; }, - Number{})); + generate_tie([&](auto i) -> const auto& // return type should be reference + { return ds_grid_desc_mblock_mperblock_nblock_nperblock[i]; }, + Number{})); // tuple of reference to C/Ds tensor descriptors const auto c_ds_buf_refs = concat_tuple_of_reference( tie(c_shuffle_block_buf), - generate_tie( - [&](auto i) -> const auto& // return type should be reference - { return ds_grid_buf[i]; }, - Number{})); + generate_tie([&](auto i) -> const auto& // return type should be reference + { return ds_grid_buf[i]; }, + Number{})); // tuple of starting index of C/Ds blockwise copy const auto idx_c_ds_block_begin = @@ -1999,41 +2000,41 @@ struct GridwiseMoeGemmMXBNS const auto EGlobalMemoryDataOperation = CGlobalMemoryDataOperation; constexpr index_t scatter_weight_idx = 3; // hack fix felix auto cde_block_copy_lds_and_global = ThreadGroupTensorSliceTransfer_v7r3_scatter< - ThisThreadBlock, - decltype(container_concat(make_tuple(CShuffleDataType{}), DsDataType{})), - Tuple, - decltype(c_ds_desc_refs), - decltype(tie(e_grid_desc_mblock_mperblock_nblock_nperblock)), - CElementwiseOperation, - Sequence(EGlobalMemoryDataOperation)>, // FIXME: make - // Sequence support - // arbitray type - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, - CDEBlockTransferCluster, - Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, - Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, - Sequence<0, 1, 2, 3>, // typename DstDimAccessOrder, - 3, // index_t SrcVectorDim, - 3, // index_t DstVectorDim, - CDEShuffleBlockTransferScalarPerVectors, - CShuffleBlockTransferScalarPerVector_NPerBlock, - sequence_merge_t< - Sequence, - uniform_sequence_gen_t>, // ThreadTransferSrcResetCoordinateAfterRunFlags - Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags - IndexType, - 1, // ScatterDim - true, // OutputScatter: false, only use scatter weights - scatter_weight_idx // ScatterWeightIdx: ascale - >{c_ds_desc_refs, - idx_c_ds_block_begin, - tie(e_grid_desc_mblock_mperblock_nblock_nperblock), - make_tuple(make_multi_index(0, 0, block_n_id, 0)), - c_element_op}; + ThisThreadBlock, + decltype(container_concat(make_tuple(CShuffleDataType{}), DsDataType{})), + Tuple, + decltype(c_ds_desc_refs), + decltype(tie(e_grid_desc_mblock_mperblock_nblock_nperblock)), + CElementwiseOperation, + Sequence(EGlobalMemoryDataOperation)>, // FIXME: make + // Sequence support + // arbitray type + Sequence<1, + CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, + CDEBlockTransferCluster, + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, + Sequence<0, 1, 2, 3>, // typename DstDimAccessOrder, + 3, // index_t SrcVectorDim, + 3, // index_t DstVectorDim, + CDEShuffleBlockTransferScalarPerVectors, + CShuffleBlockTransferScalarPerVector_NPerBlock, + sequence_merge_t< + Sequence, + uniform_sequence_gen_t>, // ThreadTransferSrcResetCoordinateAfterRunFlags + Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags + IndexType, + 1, // ScatterDim + true, // OutputScatter: false, only use scatter weights + scatter_weight_idx // ScatterWeightIdx: ascale + >{c_ds_desc_refs, + idx_c_ds_block_begin, + tie(e_grid_desc_mblock_mperblock_nblock_nperblock), + make_tuple(make_multi_index(0, 0, block_n_id, 0)), + c_element_op}; auto c_grid_buf = make_dynamic_buffer( p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); @@ -2135,6 +2136,7 @@ struct GridwiseMoeGemmMXBNS }); } } +#endif template = max_token_id || token0 >= problem.NumTokens) return; StaticallyIndexedArray gather_offsets; static_for<0, AMRepeats, 1>{}([&](auto m0) { - const index_t fused_token = p_sorted_token_ids[token_pos + m0]; + const index_t fused_token = p_sorted_token_ids[token_pos + m0 * AMThreads]; index_t token_offset = fused_token & 0xffffff; if constexpr(!IsInputGemm) { @@ -2242,11 +2244,13 @@ struct GridwiseMoeGemmMXBNS gather_offsets(m0) = static_cast(token_offset) * problem.K; }); -#if 1 - printf("blkx: %u, blky: %u, tidx: %u, token_pos: %d, gather_offsets:<%d, %d, %d, %d>\n", +#if 0 + printf("blkx: %u, blky: %u, tidx: %u,AMThreads: %d, token_pos: %d, gather_offsets:<%d, %d, " + "%d, %d>\n", blockIdx.x, blockIdx.y, threadIdx.x, + AMThreads, token_pos, gather_offsets[Number<0>{}], gather_offsets[Number<1>{}],