perf: more options for kv tile size (#336)

For small query size setting, we might use large kv tile size.
flashinfer-ai · Jun 27, 2024 · bf2a6c7 · bf2a6c7
1 parent ea89492
commit bf2a6c7
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh
@@ -72,7 +72,13 @@
   }
 
 #define DISPATCH_NUM_FRAGS_Z(max_frags_z, NUM_FRAGS_Z, ...) \
-  if (max_frags_z >= 4) {                                   \
+  if (max_frags_z >= 8) {                                   \
+    constexpr size_t NUM_FRAGS_Z = 8;                       \
+    __VA_ARGS__                                             \
+  } else if (max_frags_z >= 6) {                            \
+    constexpr size_t NUM_FRAGS_Z = 6;                       \
+    __VA_ARGS__                                             \
+  } else if (max_frags_z >= 4) {                            \
     constexpr size_t NUM_FRAGS_Z = 4;                       \
     __VA_ARGS__                                             \
   } else if (max_frags_z >= 2) {                            \