From 8b745f2c54bc0cd948a0d1009a614e9a2d1b817e Mon Sep 17 00:00:00 2001 From: Jim Date: Sat, 8 Feb 2025 17:16:24 +0800 Subject: [PATCH] codegen template --- .../ck_tile/01_fmha/codegen/cpp_symbol_map.py | 6 + .../ck_tile/01_fmha/codegen/ops/fmha_bwd.py | 226 ++++-------------- 2 files changed, 54 insertions(+), 178 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py index 332707eafd..ac02f8d8d6 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -126,3 +126,9 @@ BOOL_MAP = { "t" : "true", "f" : "false" } + +BF16_CVT_MAP = { + 0 : "rtne", + 1 : "rtna", + 2 : "rtz", +} \ No newline at end of file diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 3d0a2dd52a..6f97a65013 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -161,6 +161,13 @@ std::string fmha_bwd_dq_dk_dv_get_name_() """ FMHA_BWD_API_FILENAME="fmha_bwd_api.cpp" + +FMHA_BWD_V3_TEMPLATE=""" +template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_{F_hdim_name}_{F_dtype_name}_{F_causal_name}_{F_atomic_name}_{F_bf16_cvt_name}"; }}; +template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_{F_hdim_name}_{F_dtype_name}_{F_causal_name}_{F_atomic_name}_{F_bf16_cvt_name}; }}; +template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = {F_Ts_qo}; static constexpr int ts_kv = 192; }}; +""" + FMHA_BWD_API=""" #include #include "hsaco/fmha_hsaco.hpp" @@ -311,158 +318,7 @@ struct fmha_bwd_dq_dk_dv_v3_traits_ static constexpr bool kIsHDPad = kIsHDPad_; }}; -template struct FmhaBwdV3Name; -// ########################################################|HDim| DataType|kIsCausal|kIsAtomic32|BF16Cvt|kIsHDPad| -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_a16"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_a32"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_causal_a16"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_causal_a32"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtne_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtna_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a16_rtz_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtne_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtna_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_a32_rtz_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtne_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtna_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a16_rtz_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtne_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtna_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_bf16_causal_a32_rtz_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_a16_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_a32_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_causal_a16_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_fp16_causal_a32_pddv"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a16_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a16_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a16_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a32_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a32_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_a32_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a16_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a16_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a16_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a32_rtne"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a32_rtna"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_bf16_causal_a32_rtz"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_fp16_a16"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_fp16_a32"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_fp16_causal_a16"; }}; -template<> struct FmhaBwdV3Name> {{ static constexpr const char * bwd_v3_name = "bwd_v3_hd64_fp16_causal_a32"; }}; - -template struct FmhaBwdV3Buf; -// #######################################################|HDim| DataType|kIsCausal|kIsAtomic32|BF16Cvt|kIsHDPad| -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_a16; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_a32; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_causal_a16; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_causal_a32; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtne_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtna_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a16_rtz_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtne_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtna_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_a32_rtz_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtne_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtna_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a16_rtz_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtne_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtna_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_bf16_causal_a32_rtz_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_a16_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_a32_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_causal_a16_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_fp16_causal_a32_pddv; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a16_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a16_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a16_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a32_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a32_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_a32_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a16_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a16_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a16_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a32_rtne; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a32_rtna; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_bf16_causal_a32_rtz; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_fp16_a16; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_fp16_a32; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_fp16_causal_a16; }}; -template<> struct FmhaBwdV3Buf> {{ static constexpr unsigned char * bwd_v3_buf = bwd_hd64_fp16_causal_a32; }}; - -template struct FmhaBwdV3Ts; -// ######################################################|HDim| DataType|kIsCausal|kIsAtomic32|BF16Cvt|kIsHDPad| -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 16; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; -template<> struct FmhaBwdV3Ts> {{ static constexpr int ts_qo = 32; static constexpr int ts_kv = 192; }}; +{F_template} class fmha_bwd_v3_kernel {{ @@ -732,7 +588,17 @@ float fmha_bwd_v3_hdp_xqa_(const ck_tile::stream_config& s, fmha_bwd_args a) float fmha_bwd(fmha_bwd_traits t, fmha_bwd_args a, const ck_tile::stream_config& s){{ float r = -1; + + if ((t.uses_bwd_v3 == true){{ + if ((t.is_group_mode == false) && (t.bias_type == bias_enum::no_bias) && (t.has_dbias == false) && (t.has_dropout == false) && (t.is_deterministic == false) && (a.hdim_q == a.hdim_v) && + (a.seqlen_q == a.seqlen_k) && (a.nhead_q % a.nhead_k == 0) && (a.stride_q == a.stride_do) && (a.nhead_stride_q == a.nhead_stride_do) && (a.batch_stride_q == a.batch_stride_do) && + (a.stride_k == a.stride_v) && (a.nhead_stride_k == a.nhead_stride_v) && (a.batch_stride_k == a.batch_stride_v) && (a.nhead_stride_k == a.nhead_stride_dk) && (a.nhead_stride_v == a.nhead_stride_dv) && + ((a.batch_stride_dk / a.batch_stride_k) == (a.nhead_q / a.nhead_k)) && ((a.batch_stride_dv / a.batch_stride_v) == (a.nhead_q / a.nhead_k))) {{ + if (((a.hdim_q >= 64) && (a.hdim_q <= 128) && (a.hdim_q % 8 == 0) && (a.seqlen_k % 64 == 0))) {{ {F_v3_dispatch} + }} + }} + }} {F_dispatch} return r; }} @@ -757,23 +623,18 @@ FMHA_BWD_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) }} """ -FMHA_V3_DISPATCH=""" if ((t.uses_bwd_v3 == {F_use_v3}){{ - if ((t.is_group_mode == false) && (t.bias_type == bias_enum::no_bias) && (t.has_dbias == false) && (t.has_dropout == false) && (t.is_deterministic == false) && (a.hdim_q == a.hdim_v) && - (a.seqlen_q == a.seqlen_k) && (a.nhead_q % a.nhead_k == 0) && (a.stride_q == a.stride_do) && (a.nhead_stride_q == a.nhead_stride_do) && (a.batch_stride_q == a.batch_stride_do) && - (a.stride_k == a.stride_v) && (a.nhead_stride_k == a.nhead_stride_v) && (a.batch_stride_k == a.batch_stride_v) && (a.nhead_stride_k == a.nhead_stride_dk) && (a.nhead_stride_v == a.nhead_stride_dv) && - ((a.batch_stride_dk / a.batch_stride_k) == (a.nhead_q / a.nhead_k)) && ((a.batch_stride_dv / a.batch_stride_v) == (a.nhead_q / a.nhead_k))) {{ - if (((a.hdim_q >= 64) && (a.hdim_q <= 128) && (a.hdim_q % 8 == 0) && (a.seqlen_k % 64 == 0))) {{ +FMHA_V3_DISPATCH=""" if (t.mask_type == mask_enum::no_mask){{ if ((t.is_v3_atomic_fp32 == true) && (a.nhead_stride_dq_acc >= a.stride_dq_acc /*dq_acc only support BHSD*/)){{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}>; - using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, false, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; + using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, {F_is_causal}, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}, false>; r = fmha_bwd_v3{F_padding_suffix}_xqa_(s, a); return r; }} else if (t.is_v3_atomic_fp32 == false){{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}>; - using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, false, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; + using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, {F_is_causal}, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; r = fmha_bwd_v3{F_padding_suffix}_xqa_(s, a); return r; }} @@ -781,21 +642,19 @@ FMHA_V3_DISPATCH=""" if ((t.uses_bwd_v3 == {F_use_v3}){{ else if ((t.mask_type != mask_enum::no_mask) && ((a.window_size_left == -1) && (a.window_size_right == 0))){{ if ((t.is_v3_atomic_fp32 == true) && (a.nhead_stride_dq_acc >= a.stride_dq_acc /*dq_acc only support BHSD*/)){{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}>; - using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, true, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; + using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, {F_is_causal}, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}, false>; r = fmha_bwd_v3{F_padding_suffix}_xqa_(s, a); return r; }} else if (t.is_v3_atomic_fp32 == false){{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, false, false, {F_padding}>; - using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, true, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; + using dq_dk_dv_v3_traits_ = fmha_bwd_dq_dk_dv_v3_traits_<{F_hdim}, {F_dtype}, {F_is_causal}, {F_is_atomic32}, {F_how_v3_bf16_cvt}, {F_padding}>; r = fmha_bwd_v3{F_padding_suffix}_xqa_(s, a); return r; }} }} }} - }} - }} """ @dataclass @@ -818,9 +677,6 @@ class FmhaBwdDQDKDVApiTrait: dpad : str dvpad : str deterministic : str - uses_bwd_v3 : str - v3_atomic_fp32 : str - how_v3_bf16_cvt : int def scheck(self, spad1 : str) -> str: if self.mode == 'group': @@ -851,6 +707,15 @@ class FmhaBwdDQDKDVApiTrait: if self.dvpad == 't': return f'a.hdim_v % {self.bhdv} != 0' else : return f'a.hdim_v % {self.bhdv} == 0' +@dataclass +class FmhaBwdV3DQDKDVApiTrait: + hdim : str + dtype : str # data type + is_causal : str + is_atomic : str + bf16_cvt : int + hd_pad : str + class FmhaBwdApiPool: def __init__(self, mask_impl): self.dq_dk_dv_pool = dict() @@ -885,7 +750,6 @@ class FmhaBwdApiPool: F_scheck=trait.scheck(spad1=spad1), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_hdim=hdim, F_dtype=BWD_DTYPE_MAP[dtype], F_spad0=BOOL_MAP[trait.spad], F_spad1=BOOL_MAP[spad1], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad], F_deterministic=BOOL_MAP[trait.deterministic]) - if_j = 'if' if j == 0 else 'else if' per_hdim_case = per_hdim_case + FMHA_BWD_API_PER_HDIM_CASE.format(F_if=if_j, F_hdim=hdim, F_inner_dispatch=inners) if_i = 'if' if i == 0 else 'else if' @@ -893,20 +757,30 @@ class FmhaBwdApiPool: if not per_dtypes: # empty string we add some ignore to suppress warning in api per_dtypes += ' (void)t ; (void)s ; (void)a;' - # v3 codegen: + v3_code = str() + gen_template = str() for i, dtype in enumerate(self.dq_dk_dv_pool.keys()): - per_hdim_case=str() for j, hdim in enumerate(self.dq_dk_dv_pool[dtype].keys()): traits = self.dq_dk_dv_pool[dtype][hdim] hdim_int = int(hdim) hdim_int = (hdim_int + 64 - 1) / 64 * 64 + Ts_qo = 32 if hdim == 64 else 16 for k, trait in enumerate(traits): padding = "t" if hdim_int % 64 == 0 else "f" padding_suffix = "_hdp" if padding == "t" else "" - v3_code = v3_code + FMHA_V3_DISPATCH.format(F_use_v3=BOOL_MAP[trait.use_bwd_v3], F_hdim=hdim_int, F_dtype=BWD_DTYPE_MAP[dtype], F_padding=BOOL_MAP[padding], - F_is_atomic32=BOOL_MAP[trait.v3_atomic_fp32], F_how_v3_bf16_cvt=trait.how_v3_bf16_cvt, F_padding_suffix=padding_suffix) - return FMHA_BWD_KERNEL_HEADER + FMHA_BWD_API.format(F_dispatch = per_dtypes, F_v3_dispatch = v3_code) + v3_code = v3_code + FMHA_V3_DISPATCH.format(F_hdim=hdim_int, F_dtype=BWD_DTYPE_MAP[dtype], F_padding=BOOL_MAP[padding], + F_is_atomic32=BOOL_MAP[trait.is_atomic], F_how_v3_bf16_cvt=trait.bf16_cvt, F_padding_suffix=padding_suffix) + + hdim_name = "hd64" if hdim_int == 64 else "" + dtype_name = dtype + causal_name = "causal" if trait.is_causal == "t" else "" + atomic_name = "a32" if trait.is_atomic == "t" else "a16" + bf16_cvt_name = BF16_CVT_MAP[trait.bf16_cvt] + gen_template = gen_template + FMHA_BWD_V3_TEMPLATE.format(F_hdim=hdim_int, F_dtype=BWD_DTYPE_MAP[dtype], F_is_atomic=BOOL_MAP[trait.is_atomic], + F_is_causal=BOOL_MAP[trait.is_causal], F_bf16_cvt=trait.bf16_cvt, F_hdpad=BOOL_MAP[padding], F_Ts_qo = Ts_qo, F_hdim_name=hdim_name, + F_dtype_name=dtype_name, F_causal_name=causal_name, F_atomic_name=atomic_name, F_bf16_cvt_name=bf16_cvt_name) + return FMHA_BWD_KERNEL_HEADER + FMHA_BWD_API.format(F_dispatch = per_dtypes, F_template = gen_template, F_v3_dispatch = v3_code) # GEMM0: Q@K=S^T # GEMM1: P^T@dO^T=dV(This was chosen as G1 to match fwd, but N1 must be equal to headdim_v) @@ -1038,7 +912,6 @@ class FmhaBwdDQDKDVKernel: def filename(self) -> str: return self.name + ".cpp" - # TODO: Jim: update this def api_trait(self) -> FmhaBwdDQDKDVApiTrait: return FmhaBwdDQDKDVApiTrait(pipeline=self.F_pipeline, hdim=str(self.F_hdim), @@ -1056,10 +929,7 @@ class FmhaBwdDQDKDVKernel: skpad=self.F_skpad, dpad=self.F_dpad, dvpad=self.F_dvpad, - deterministic=self.F_deterministic, - use_bwd_v3=self.use_bwd_v3, - v3_atomic_fp32=self.v3_atomic_fp32, - how_v3_bf16_cvt=self.how_v3_bf16_cvt + deterministic=self.F_deterministic ) # TODO: design a more practical way to do it