Revert "Try to have kPadSeqLenK be false in splitkv dispatch"

This reverts commit 7243b49.
ROCm · Oct 26, 2024 · 6ffea6a · 6ffea6a
1 parent 7243b49
commit 6ffea6a
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 12 deletions.
diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward_splitkv_dispatch.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward_splitkv_dispatch.h
@@ -96,9 +96,7 @@ struct batched_forward_splitkv_causalmask_bias_dropout_dispatch {
             has_uneven_splits,
             kHasUnevenSplits,
             [&] {
-              // since buffer_load_dword is used, padding dim seqlen-k is not
-              // needed when loading K/V, but still needed when loading bias
-              constexpr bool kPadSeqLenK = kHasBias ? true : false;
+              constexpr bool kPadSeqLenK = kHasUnevenSplits ? true : false;
 
               using FmhaTraits = ck_tile::TileFmhaFwdSplitKVTraits<
                   kPadSeqLenQ,

diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_splitkv_dispatch.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_splitkv_dispatch.h
@@ -96,9 +96,7 @@ struct batched_infer_splitkv_causalmask_bias_dropout_dispatch {
             has_uneven_splits,
             kHasUnevenSplits,
             [&] {
-              // since buffer_load_dword is used, padding dim seqlen-k is not
-              // needed when loading K/V, but still needed when loading bias
-              constexpr bool kPadSeqLenK = kHasBias ? true : false;
+              constexpr bool kPadSeqLenK = kHasUnevenSplits ? true : false;
 
               using FmhaTraits = ck_tile::TileFmhaFwdSplitKVTraits<
                   kPadSeqLenQ,

diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_forward_splitkv_dispatch.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_forward_splitkv_dispatch.h
@@ -78,9 +78,7 @@ struct grouped_forward_splitkv_causalmask_bias_dropout_dispatch {
             : ck_tile::BlockAttentionBiasEnum::NO_BIAS;
 
         constexpr bool kPadSeqLenQ = true;
-        // since buffer_load_dword is used, padding dim seqlen-k is not
-        // needed when loading K/V, but still needed when loading bias
-        constexpr bool kPadSeqLenK = kHasBias? true : false;
+        constexpr bool kPadSeqLenK = true;
 
         const bool pad_headdim_q =
             !(param.K % FmhaTileShape::kK0BlockLength == 0);

diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_infer_splitkv_dispatch.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_infer_splitkv_dispatch.h
@@ -78,9 +78,7 @@ struct grouped_infer_splitkv_causalmask_bias_dropout_dispatch {
             : ck_tile::BlockAttentionBiasEnum::NO_BIAS;
 
         constexpr bool kPadSeqLenQ = true;
-        // since buffer_load_dword is used, padding dim seqlen-k is not
-        // needed when loading K/V, but still needed when loading bias
-        constexpr bool kPadSeqLenK = kHasBias ? true : false;
+        constexpr bool kPadSeqLenK = true;
 
         bool pad_headdim_q = !(param.K % FmhaTileShape::kK0BlockLength == 0);
         bool pad_headdim_v = !(param.Kv % FmhaTileShape::kN1 == 0);