From 9ec586fc6884fcff4b91a909fb0f866001a5cb39 Mon Sep 17 00:00:00 2001 From: letaoqin Date: Wed, 27 Nov 2024 09:43:04 +0000 Subject: [PATCH] change a matrix lds desc --- .../instances/fused_moegemm_api.cpp | 2 +- .../instances/fused_moegemm_bf16_m32.cpp | 2 +- .../fused_moegemm_pipeline_general_policy.hpp | 98 +++---------------- 3 files changed, 18 insertions(+), 84 deletions(-) diff --git a/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_api.cpp b/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_api.cpp index 8823835aa4..7aa6185b8c 100644 --- a/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_api.cpp +++ b/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_api.cpp @@ -19,7 +19,7 @@ float fused_moegemm(fused_moegemm_traits t, fused_moegemm_args a, const ck_tile: if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" && t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1) { - using t_ = fmoe_, S<1, 4, 1>, S<32, 32, 8>, 1, 0>; + using t_ = fmoe_, S<1, 4, 1>, S<32, 32, 8>, 1, 0>; r = fused_moegemm_(s, a); } // clang-format on diff --git a/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_bf16_m32.cpp b/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_bf16_m32.cpp index a4dc3693b6..5fd456c1c8 100644 --- a/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_bf16_m32.cpp +++ b/example/ck_tile/16_fused_moe_general/instances/fused_moegemm_bf16_m32.cpp @@ -8,7 +8,7 @@ // clang-format off template float fused_moegemm_< - fmoe_, S<1, 4, 1>, S<32, 32, 8>, 1, 0> + fmoe_, S<1, 4, 1>, S<32, 32, 8>, 1, 0> >(const ck_tile::stream_config& s, fused_moegemm_args a); // clang-format on diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general_policy.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general_policy.hpp index d301f884af..ac65e9b975 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general_policy.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general_policy.hpp @@ -201,7 +201,7 @@ struct FusedMoeGemmPipelineGeneralPolicy using S_ = typename Problem::BlockShape; constexpr index_t K2 = S_::Warp_K0; constexpr index_t K1 = get_warp_size() / S_::Warp_N0; - constexpr index_t K0 = kKPerBlock / (K1 * K2); + constexpr index_t K0 = S_::Repeat_K0; return make_static_tile_distribution( tile_distribution_encoding< @@ -277,93 +277,27 @@ struct FusedMoeGemmPipelineGeneralPolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeLdsStoreDesc_A() { - // A async->LDS constexpr index_t Block_M = Problem::BlockShape::Block_M0; constexpr index_t Block_K = Problem::BlockShape::Block_K0; - // constexpr index_t BlockSize = Problem::BlockShape::BlockSize; - constexpr index_t warpSize = ck_tile::get_warp_size(); - constexpr index_t NumWarps = Problem::BlockShape::NumWarps; + constexpr index_t kK1 = GetSmemKPack_A(); // LDS + constexpr index_t kK0 = Block_K / kK1; - constexpr index_t KPack = GetSmemKPack_A(); // LDS - constexpr index_t KVector = GetAlignment_A(); // async copy 1 dword - constexpr index_t KPad = KPack; // pad between warps + static_assert(Block_K % kK1 == 0); - static_assert(Block_K % KVector == 0); - constexpr index_t LanesPerK = Block_K / KVector; // how many thread loading K - if constexpr(LanesPerK >= warpSize) - { - // need multiple waves to load K - static_assert(LanesPerK % warpSize == 0); - constexpr index_t wavesPerK = LanesPerK / warpSize; - if constexpr(wavesPerK > NumWarps) - { - // TODO: need multiple issues along K to load all data - } - else - { - constexpr index_t wavesPerM = NumWarps / wavesPerK; - constexpr index_t NumIssues = Block_M / wavesPerM; - constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, // m0 - number{}, // m1 - number{}, // k0 - number{}, // k1 - number{}), // k2 - make_tuple(number{}, // m0 - number{}, // m1 - number{}, // k0 - number{}, // k1 - number<1>{}), // k2 - number{}, // lds store vector(actually no explicit store) - number<1>{}); + constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, number{}, number{}), + make_tuple(number<(Block_M + 1) * kK1>{}, number{}, number<1>{}), + number<8>{}, + number<1>{}); - constexpr auto lds_block_desc_issues_warps_lanes = transform_tensor_descriptor( - lds_block_desc_0, - make_tuple( - // make_pass_through_transform(), - make_merge_transform(make_tuple(number{}, number{})), - make_merge_transform(make_tuple( - number{}, number{}, number{}))), - make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}), - make_tuple(sequence<0>{}, sequence<1>{})); + constexpr auto a_lds_block_desc = transform_tensor_descriptor( + a_lds_block_desc_0, + make_tuple(make_pass_through_transform(number{}), + make_merge_transform(make_tuple(number{}, number{}))), + make_tuple(sequence<1>{}, sequence<0, 2>{}), + make_tuple(sequence<0>{}, sequence<1>{})); - return lds_block_desc_issues_warps_lanes; - } - } - else - { - // lanes within a wave load different M but same K - static_assert(warpSize % LanesPerK == 0); - constexpr index_t LaneGroups = warpSize / LanesPerK; // along m - constexpr index_t NumIssues = Block_M / (LaneGroups * NumWarps); - - constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, // m0 - number{}, // m1 - number{}, // m2 - number{}, // k0 - number{}), // k1 - make_tuple(number{}, // m0 - number{}, // m1 - number{}, // m2 - number{}, // k0 - number<1>{}), // k1 - number{}, // lds store vector(actually no explicit store) - number<1>{}); - - constexpr auto lds_block_desc_issues_warps_lanes = transform_tensor_descriptor( - lds_block_desc_0, - make_tuple( - // make_pass_through_transform(number{}), - // make_pass_through_transform(number{}), - make_merge_transform( - make_tuple(number{}, number{}, number{})), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<0, 1, 2>{}, sequence<3, 4>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - - return lds_block_desc_issues_warps_lanes; - } + return a_lds_block_desc; } template