formatting code

2026-07-19 02:01:01 +00:00 · 2025-08-14 02:01:06 +08:00
parent 777c6e8f53
commit 178e0b71c9
11 changed files with 60 additions and 104 deletions
--- a/example/ck_tile/01_fmha/CMakeLists.txt
+++ b/example/ck_tile/01_fmha/CMakeLists.txt
@@ -138,7 +138,6 @@ endif()
 # Allow comparing floating points directly in order to check sentinel values
 list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-float-equal)
 list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-float-equal)
-# list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS --save-temps)

 target_compile_options(${EXAMPLE_FMHA_FWD} PRIVATE ${EXAMPLE_FMHA_FWD_COMPILE_OPTIONS})
 target_compile_options(${EXAMPLE_FMHA_BWD} PRIVATE ${EXAMPLE_FMHA_BWD_COMPILE_OPTIONS})
--- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
+++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
@@ -286,11 +286,8 @@ class FmhaBwdDQDKDVApiTrait:
        else :                         return '!t.is_deterministic && !t.is_atomic_fp32'

    def get_kernel_group(self) -> Tuple[str]:
-        # kernel_group = [FMHA_BWD_CONVERT_DQ]
        kernel_group = [FMHA_BWD_DOT_DO_O, FMHA_BWD_DQ_DQ_DV, FMHA_BWD_CONVERT_DQ]
-        # if self.deterministic == 't' or self.atomic32 == 't':
-        #     kernel_group.append(FMHA_BWD_CONVERT_DQ)
-        return tuple(kernel_group)  
+        return tuple(kernel_group)

 class FmhaBwdApiPool:
    def __init__(self, mask_impl):
@@ -568,10 +565,10 @@ class FmhaBwdDQDKDVKernel:
 def get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype : str) -> Optional[dict]:
    if dtype == 'fp16' or dtype == 'bf16':
        return {
-            # '32'  : [FmhaBwdDQDKDVTileSize( 32, 128,  32, 32,  32, 32, 64,  32,  32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
-            #             "kr_ktr_vr_iglp", "kr_ktr_vr"],
-            # '64'  : [FmhaBwdDQDKDVTileSize( 32, 128,  64, 32,  64, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
-            #             "kr_ktr_vr_iglp", "kr_ktr_vr"],
+            '32'  : [FmhaBwdDQDKDVTileSize( 32, 128,  32, 32,  32, 32, 64,  32,  32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
+            '64'  : [FmhaBwdDQDKDVTileSize( 32, 128,  64, 32,  64, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
            '128' : [FmhaBwdDQDKDVTileSize( 16, 128, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
            '256' : [FmhaBwdDQDKDVTileSize( 16,  64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
@@ -590,9 +587,6 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
        if d == None:
            continue
        for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad, deterministic, atomic32 in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], DROPOUT_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]):
-            # for debug(xiangxli)
-            if bias != 'no' or dropout!= 'no' or deterministic == 't':
-                continue

            tile = d[hdim_str][0]
            ppl = d[hdim_str][1]
--- a/example/ck_tile/01_fmha/fmha_bwd.cpp
+++ b/example/ck_tile/01_fmha/fmha_bwd.cpp
@@ -95,7 +95,7 @@ auto create_args(int argc, char* argv[])
                "0",
                "if set to 1 will use multi-buffer reduction strategy for dq, atomic opeartion "
                "will not be used")
-        .insert("atomic_fp32", "1", "if set to 0 will use atomic fp16/bf16(w/o convert_dq kernel)");
+        .insert("atomic_fp32", "1", "if set to 0 will use atomic fp16/bf16");

    bool result = arg_parser.parse(argc, argv);
    return std::make_tuple(result, arg_parser);
@@ -123,7 +123,7 @@ auto get_elimit<FmhaBwdBf16>(ck_tile::index_t hdim_q, ck_tile::index_t hdim_v)
    return ck_tile::make_tuple(rtol, atol);
 }

-ck_tile::index_t get_bit_ceil(const ck_tile::index_t& dim_value)
+ck_tile::index_t get_bit_ceil(const ck_tile::index_t dim_value)
 {
    unsigned un = static_cast<unsigned>(dim_value);
    un |= un >> 1;
@@ -215,9 +215,8 @@ bool run(const ck_tile::ArgParser& arg_parser)

    // for dq_acc padding in atomic16
    constexpr ck_tile::index_t seqlen_dq_acc_tile_size = 16;
-    const ck_tile::index_t hdim_q_pad = get_bit_ceil(hdim_q);
-    const ck_tile::index_t hdim_q_dq_acc = atomic_fp32 ? hdim_q : hdim_q_pad;
-
+    const ck_tile::index_t hdim_q_pad                  = get_bit_ceil(hdim_q);
+    const ck_tile::index_t hdim_q_dq_acc               = atomic_fp32 ? hdim_q : hdim_q_pad;

    ck_tile::stream_config stream_config{nullptr,
                                         true,
@@ -230,10 +229,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
    const auto seqstart_k_host = generate_seqstarts(mode, batch, seqlen_k);

    auto seqstart_dq_acc_host = std::vector<int32_t>(seqstart_q_host.size(), 0);
-    for (int i = 0; i < batch; ++i) {
-        auto cur_seqlen_q = seqstart_q_host[i+1] - seqstart_q_host[i];
-        auto cur_seqlen_dq_acc = ck_tile::integer_least_multiple(cur_seqlen_q, seqlen_dq_acc_tile_size);
-        seqstart_dq_acc_host[i+1] = seqstart_dq_acc_host[i] + cur_seqlen_dq_acc;
+    for(int i = 0; i < batch; ++i)
+    {
+        auto cur_seqlen_q = seqstart_q_host[i + 1] - seqstart_q_host[i];
+        auto cur_seqlen_dq_acc =
+            ck_tile::integer_least_multiple(cur_seqlen_q, seqlen_dq_acc_tile_size);
+        seqstart_dq_acc_host[i + 1] = seqstart_dq_acc_host[i] + cur_seqlen_dq_acc;
    }

    using TypeConfig = FmhaBwdTypeConfig<DataTypeConfig>;
@@ -524,7 +525,6 @@ bool run(const ck_tile::ArgParser& arg_parser)
        const ck_tile::index_t batch_stride_dq_acc  = (nhead * shape_seqlen_dq_acc * hdim_q_dq_acc);
        const ck_tile::index_t split_stride_dq_acc =
            (shape_batch * nhead * shape_seqlen_dq_acc * hdim_q_dq_acc);
-        auto dq_acc_ptr = dq_acc_buf.GetDeviceBuffer();
        const auto drop_seed_offset = [&]() -> decltype(fmha_bwd_args::drop_seed_offset) {
            if(drop_prefs)
            {
@@ -551,7 +551,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                             dk_buf.GetDeviceBuffer(),
                             dv_buf.GetDeviceBuffer(),
                             dbias_buf.GetDeviceBuffer(),
-                             dq_acc_ptr,
+                             dq_acc_buf.GetDeviceBuffer(),
                             seqstart_q.GetDeviceBuffer(),
                             seqstart_k.GetDeviceBuffer(),
                             nullptr,
@@ -1022,22 +1022,26 @@ int main(int argc, char* argv[])
        return -1;

    const std::string data_type = arg_parser.get_str("prec");
-    const bool atomic_fp32   = arg_parser.get_bool("atomic_fp32");
+    const bool atomic_fp32      = arg_parser.get_bool("atomic_fp32");
    if(data_type == "fp16")
    {
-        if (atomic_fp32) {
+        if(atomic_fp32)
+        {
            return run<FmhaBwdFp16>(arg_parser) ? 0 : -2;
        }
-        else {
+        else
+        {
            return run<FmhaBwdFp16, false>(arg_parser) ? 0 : -2;
        }
    }
    else if(data_type == "bf16")
    {
-        if (atomic_fp32) {
+        if(atomic_fp32)
+        {
            return run<FmhaBwdBf16>(arg_parser) ? 0 : -2;
        }
-        else {
+        else
+        {
            return run<FmhaBwdBf16, false>(arg_parser) ? 0 : -2;
        }
    }
--- a/example/ck_tile/01_fmha/fmha_bwd.hpp
+++ b/example/ck_tile/01_fmha/fmha_bwd.hpp
@@ -423,7 +423,7 @@ template <ck_tile::index_t HDim_,
          bool kPadS_,
          bool kPadD_,
          bool kIsDeterministic_,
-          bool kAtomic32_>
+          bool kAtomic32_ = true>
 struct fmha_bwd_convert_dq_traits_
 {
    static constexpr ck_tile::index_t HDim = HDim_;
--- a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh
+++ b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh
@@ -16,13 +16,14 @@ for bias in "n" "a" ; do
 for dbias in 0 ; do
 for p_drop in 0.0 0.2 ; do
 for deterministic in 0 ; do
+for atomic_fp32 in 0 1 ; do

-$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS

 done
 done
@@ -32,4 +33,5 @@ done
 done
 done
 done
+done
 set +x