diff --git a/example/ck_tile/38_block_scale_gemm/README.md b/example/ck_tile/38_block_scale_gemm/README.md index 496697ca32..c69703e50c 100644 --- a/example/ck_tile/38_block_scale_gemm/README.md +++ b/example/ck_tile/38_block_scale_gemm/README.md @@ -64,8 +64,10 @@ User need to select correct mapping of config for each quant mode: | | quant_mode as runtime argument | Config in cpp file | |:--------|:-----:|-------| | For selecting AQuant | aquant | GemmConfigQuant | -| For selecting Aquant with Preshuffle | aquant | GemmConfigPreshuffleQuant | +| For selecting AQuant with Preshuffle | aquant | GemmConfigPreshuffleQuant | | For selecting BQuant | bquant | GemmConfigQuant | -| For selecting PreShuffle Weight matrix with Bquant | bquant | GemmConfigPreshuffleB_Bquant_decode (or) GemmConfigPreshuffleB_Bquant_prefill +| For selecting BQuant with Preshuffle | bquant | GemmConfigPreshuffleQuant | +| For selecting PreShuffle B with BQuant | bquant | GemmConfigPreshuffleB_BQuant_Decode (or) GemmConfigPreshuffleB_BQuant_Prefill +| For selecting PreShuffle B with preshuffle BQuant | bquant | GemmConfigPreshuffleB_PreshuffleBQuant_Decode (or) GemmConfigPreshuffleB_PreshuffleBQuant_Prefill | For selecting RowCol quant | rowcolquant | GemmConfigRowColQuant | diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp index 63f76c5459..46cb017335 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp @@ -460,5 +460,5 @@ int run_gemm_example(int argc, char* argv[]) int main(int argc, char* argv[]) { - return !run_gemm_example(argc, argv); + return !run_gemm_example(argc, argv); } diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp index 8322efd006..aed6d40df8 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp @@ -148,7 +148,30 @@ struct GemmConfigPreshuffleQuant : public GemmConfigBase }; template -struct GemmConfigPreshuffleB_Bquant_decode : public GemmConfigBase +struct GemmConfigPreshuffleB_BQuant_Decode : public GemmConfigBase +{ + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = + get_k_from_preshuffled_warp_tile(); + + static constexpr bool PreshuffleB = true; + static constexpr bool DoubleSmemBuffer = true; + + static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; + static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; +}; + +template +struct GemmConfigPreshuffleB_PreshuffleBQuant_Decode : public GemmConfigBase { static constexpr ck_tile::index_t M_Tile = 16; static constexpr ck_tile::index_t N_Tile = 64; @@ -166,13 +189,10 @@ struct GemmConfigPreshuffleB_Bquant_decode : public GemmConfigBase static constexpr bool PreshuffleB = true; static constexpr bool DoubleSmemBuffer = true; static constexpr bool PreshuffleQuant = true; - - static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; - static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; }; template -struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase +struct GemmConfigPreshuffleB_BQuant_Prefill : public GemmConfigBase { static constexpr ck_tile::index_t M_Tile = 128; static constexpr ck_tile::index_t N_Tile = 128; @@ -194,6 +214,27 @@ struct GemmConfigPreshuffleB_Bquant_prefill : public GemmConfigBase static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; }; +template +struct GemmConfigPreshuffleB_PreshuffleBQuant_Prefill : public GemmConfigBase +{ + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = + get_k_from_preshuffled_warp_tile(); + + static constexpr bool PreshuffleB = true; + static constexpr bool DoubleSmemBuffer = true; + static constexpr bool PreshuffleQuant = true; +}; + template