pipeline fixes for accuracy issues; disable pre-softmax function until its accuracy is fixed

example/ck_tile/18_flexattn/CMakeLists.txt

@@ -8,9 +8,11 @@ endif()
 
 variable_watch(FMHA_SCORE_MOD_F)
 set(FMHA_SCORE_MOD_F [[s + static_cast<decltype(s)>((q_idx - v_idx) % 8)]])
+# set(FMHA_SCORE_MOD_F [[s]])
 
 variable_watch(FMHA_PRE_SOFTMAX_F)
-set(FMHA_PRE_SOFTMAX_F [[static_cast<decltype(s)>(tanh(s*1.0)/1.0)]])
+# set(FMHA_PRE_SOFTMAX_F [[static_cast<decltype(s)>(tanh(s*1.0)/1.0)]])
+set(FMHA_PRE_SOFTMAX_F [[s]])
 
 foreach(api ${FMHA_FWD_ENABLE_APIS})
   if(NOT "${api}" IN_LIST FMHA_FWD_KNOWN_APIS)

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_flex_qr_ks_vs.hpp

@@ -350,7 +350,7 @@ struct BlockFmhaPipelineQRKSVS
                             s_acc.get_tile_distribution(), make_tuple(idx0, idx1));
 
                         const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<1>{}) + tile_idx.at(number<1>{});
+                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
                         constexpr auto i_j_idx = make_tuple(idx0, idx1);
                         s_acc(i_j_idx)         = score_mod(s_acc(i_j_idx), row, col);
                     });

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_flex_qr_ks_vs_async.hpp

@@ -420,7 +420,7 @@ struct BlockFmhaPipelineQRKSVSAsync
                             s_acc.get_tile_distribution(), make_tuple(idx0, idx1));
 
                         const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<1>{}) + tile_idx.at(number<1>{});
+                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
                         constexpr auto i_j_idx = make_tuple(idx0, idx1);
 
                         s_acc(i_j_idx) = score_mod(s_acc(i_j_idx), row, col);