From a61084f4323250d75cfabe780cc322945c2e7f78 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Tue, 18 Feb 2025 07:04:22 +0000 Subject: [PATCH] opt gemm2 to 2x2 wave --- example/65_gemm_multiply_multiply/moe_gemm2.cpp | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/example/65_gemm_multiply_multiply/moe_gemm2.cpp b/example/65_gemm_multiply_multiply/moe_gemm2.cpp index c305d6e004..94275db226 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm2.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm2.cpp @@ -127,12 +127,14 @@ using CDEElementOp = MulABScaleExpertWeight; static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; static constexpr ck::index_t MPerBlock = 128; static constexpr ck::index_t BLOCKSIZE = 256; +static constexpr ck::index_t MXDLPerWave = 2; +static constexpr ck::index_t NXDLPerWave = 2; static constexpr ck::index_t NPerBlock = 128; static constexpr ck::index_t MNPerXDL = 32; -static constexpr ck::index_t KPerBlock = 256 / sizeof(A0DataType); -static constexpr ck::index_t MXDLPerWave = MPerBlock / 32; //todo fix this constraint -static constexpr ck::index_t CShuffleMXDLPerWave = MPerBlock / 32; -static constexpr ck::index_t CShuffleNLane = NPerBlock / 2; +static constexpr ck::index_t KPerBlock = 128 / sizeof(A0DataType); +// static constexpr ck::index_t MXDLPerWave = MPerBlock / 32; //todo fix this constraint +// static constexpr ck::index_t CShuffleMXDLPerWave = MPerBlock / 32; +static constexpr ck::index_t CShuffleNLane = 32; static constexpr ck::index_t CShuffleMLane = BLOCKSIZE / CShuffleNLane; static constexpr ck::index_t AK1 = 16 / sizeof(A0DataType); static constexpr ck::index_t BK1 = 16 / sizeof(B0DataType); @@ -159,7 +161,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm // mn_perxdl MNPerXDL, MNPerXDL, // mn_xdlperwave - MXDLPerWave, 1, + MXDLPerWave, NXDLPerWave, // a,b: loadtranfer cluster, cluster order, srcorder,VECDIM, srcpervec, dstpervec, lds_extra // S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, // S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, @@ -168,7 +170,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm // CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| // MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| // PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| - CShuffleMXDLPerWave, 1, S<1, CShuffleMLane, 1, CShuffleNLane>, S, + MXDLPerWave, 1, S<1, CShuffleMLane, 1, CShuffleNLane>, S, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, true, false, A0DataType>; // kernel 2: 128->32x128x128 // < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 128, 32, 128, 128, 16, 16, 32, 32, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 1, 1, S<1, 16, 1, 8>, S<8, 8, 1>, ck::BlockGemmPipelineScheduler::Interwave, ck::BlockGemmPipelineVersion::v1, EDataType>;