Add kMaxSplits=8 support

ROCm · Oct 17, 2024 · bc107ad · bc107ad
1 parent 4b4327e
commit bc107ad
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_num_kv_split_switch.h b/xformers/csrc/attention/hip_fmha/ck_tiled_fmha_num_kv_split_switch.h
@@ -11,7 +11,10 @@
 
 #define FMHA_FWD_NUM_KV_SPLITS_SWITCH(NUM_SPLITS, CONST_NAME, ...) \
   [&] {                                                            \
-    if (NUM_SPLITS <= 16) {                                        \
+    if (NUM_SPLITS <= 8) {                                         \
+      constexpr ck_tile::index_t CONST_NAME = 3;                   \
+      __VA_ARGS__();                                               \
+    } else if (NUM_SPLITS <= 16) {                                 \
       constexpr ck_tile::index_t CONST_NAME = 4;                   \
       __VA_ARGS__();                                               \
     } else if (NUM_SPLITS <= 32) {                                 \