Remove using splitkv from fmha-fwd training path

ROCm · Oct 16, 2024 · f94fdfd · f94fdfd
1 parent d4437ad
commit f94fdfd
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 52 deletions.
diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward.h
@@ -7,7 +7,6 @@
 #pragma once
 
 #include "ck_tiled_fmha_batched_forward_dispatch.h"
-#include "ck_tiled_fmha_batched_forward_splitkv_dispatch.h"
 
 template <
     typename ScalarType,
@@ -18,29 +17,10 @@ template <
 void run_batched_forward_causalmask_bias_dropout_dispatch(
     BatchedForwardParams& param,
     hipStream_t stream) {
-  // currently split-kv implementation does not support dropout
-  if constexpr (!kHasDropout) {
-#ifndef FMHA_FWD_SPLITKV_NOT_USED
-    if (param.use_split_kv)
-      batched_forward_splitkv_causalmask_bias_dropout_dispatch<
-          ScalarType,
-          kHasCausalMask,
-          kHasBias,
-          MaxK>::Run(param, stream);
-    else
-#endif
-      batched_forward_causalmask_bias_dropout_dispatch<
-          ScalarType,
-          kHasCausalMask,
-          kHasBias,
-          kHasDropout,
-          MaxK>::Run(param, stream);
-  } else {
-    batched_forward_causalmask_bias_dropout_dispatch<
-        ScalarType,
-        kHasCausalMask,
-        kHasBias,
-        kHasDropout,
-        MaxK>::Run(param, stream);
-  }
+  batched_forward_causalmask_bias_dropout_dispatch<
+      ScalarType,
+      kHasCausalMask,
+      kHasBias,
+      kHasDropout,
+      MaxK>::Run(param, stream);
 };
diff --git a/...d_fmha_batched_forward_splitkv_dispatch.h → ...ched_forward_splitkv_dispatch_discarded.h b/...d_fmha_batched_forward_splitkv_dispatch.h → ...ched_forward_splitkv_dispatch_discarded.h
diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_forward.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_forward.h
@@ -7,7 +7,6 @@
 #pragma once
 
 #include "ck_tiled_fmha_grouped_forward_dispatch.h"
-#include "ck_tiled_fmha_grouped_forward_splitkv_dispatch.h"
 
 template <
     typename ScalarType,
@@ -18,29 +17,10 @@ template <
 void run_grouped_forward_causalmask_bias_dropout_dispatch(
     GroupedForwardParams& param,
     hipStream_t stream) {
-  // currently split-kv implementation does not support dropout
-  if constexpr (!kHasDropout) {
-#ifndef FMHA_FWD_SPLITKV_NOT_USED
-    if (param.use_split_kv)
-      grouped_forward_splitkv_causalmask_bias_dropout_dispatch<
-          ScalarType,
-          kHasCausalMask,
-          kHasBias,
-          MaxK>::Run(param, stream);
-    else
-#endif
-      grouped_forward_causalmask_bias_dropout_dispatch<
-          ScalarType,
-          kHasCausalMask,
-          kHasBias,
-          kHasDropout,
-          MaxK>::Run(param, stream);
-  } else {
-    grouped_forward_causalmask_bias_dropout_dispatch<
-        ScalarType,
-        kHasCausalMask,
-        kHasBias,
-        kHasDropout,
-        MaxK>::Run(param, stream);
-  }
+  grouped_forward_causalmask_bias_dropout_dispatch<
+      ScalarType,
+      kHasCausalMask,
+      kHasBias,
+      kHasDropout,
+      MaxK>::Run(param, stream);
 };
diff --git a/...d_fmha_grouped_forward_splitkv_dispatch.h → ...uped_forward_splitkv_dispatch_discarded.h b/...d_fmha_grouped_forward_splitkv_dispatch.h → ...uped_forward_splitkv_dispatch_discarded.h