[AMD ROCm] Fix intrinsic for ROCm7 (Dao-AILab#1729)

* Use more reasonable splitkv heuristic

* update CK

* Pass logits soft-capping arguments

* Revert "Merge pull request #147 from ROCm/poyenc/fix-ck-tile-splitkv-heuristic"

This reverts commit 12857ce, reversing
changes made to e64b970.

---------

Co-authored-by: Po Yen Chen <[email protected]>

Author: rocking5566

csrc/composable_kernel

@@ -19,6 +19,7 @@ fmha_fwd_traits get_ck_fmha_fwd_traits(const mask_info &mask,
                            dtype,
                            false, // is_group_mode
                            true,  // is_v_rowmajor
+                           false, // has_logits_soft_cap
                            mask.type,
                            enable_alibi ? bias_enum::alibi : bias_enum::no_bias,
                            has_lse,
@@ -111,6 +112,7 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
                          softmax_scale, // scale_s
                          1,             // scale_p
                          1,             // scale_o
+                         0.0f,          // logits_soft_cap
                          stride_q,
                          stride_k,
                          stride_v,
@@ -134,6 +136,7 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
                          mask.left,
                          mask.right,
                          static_cast<ck_tile::index_t>(mask.type),
+                         0, // min_seqlen_q
                          p_dropout,
                          has_dropout_randval,
                          drop_seed_offset};

csrc/flash_attn_ck/mha_fwd.cpp

@@ -33,7 +33,8 @@ fmha_fwd_splitkv_traits get_ck_fmha_fwd_splitkv_traits(const mask_info &mask,
                                    head_size,
                                    dtype,
                                    false, // is_group_mode
-                                   true, // is_v_rowmajor
+                                   true,  // is_v_rowmajor
+                                   false, // has_logits_soft_cap
                                    mask.type,
                                    enable_alibi ? bias_enum::alibi : bias_enum::no_bias,
                                    has_lse,

csrc/flash_attn_ck/mha_fwd_kvcache.cpp

@@ -17,8 +17,9 @@ fmha_fwd_traits get_ck_fmha_varlen_fwd_traits(const mask_info &mask,
     return fmha_fwd_traits{head_size,
                            head_size,
                            dtype,
-                           true, // is_group_mode
-                           true, // is_v_rowmajor
+                           true,  // is_group_mode
+                           true,  // is_v_rowmajor
+                           false, // has_logits_soft_cap
                            mask.type,
                            enable_alibi ? bias_enum::alibi : bias_enum::no_bias,
                            has_lse,
@@ -35,8 +36,9 @@ fmha_fwd_splitkv_traits get_ck_fmha_varlen_fwd_splitkv_traits(const mask_info &m
     return fmha_fwd_splitkv_traits{head_size,
                                    head_size,
                                    dtype,
-                                   true, // is_group_mode
-                                   true, // is_v_rowmajor
+                                   true,  // is_group_mode
+                                   true,  // is_v_rowmajor
+                                   false, // has_logits_soft_cap
                                    mask.type,
                                    enable_alibi ? bias_enum::alibi : bias_enum::no_bias,
                                    has_lse,
@@ -131,6 +133,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
                          softmax_scale, // scale_s
                          1,             // scale_p
                          1,             // scale_o
+                         0.0f,          // logits_soft_cap
                          stride_q,
                          stride_k,
                          stride_v,
@@ -154,6 +157,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
                          mask.left,
                          mask.right,
                          static_cast<ck_tile::index_t>(mask.type),
+                         0, // min_seqlen_q
                          p_dropout,
                          has_dropout_randval,
                          drop_seed_offset};