feat: add group size 3 to GQA decode dispatch (#558)

Llama 3.2 3B comes with 24 qo heads and 8 kv heads.
flashinfer-ai · Oct 25, 2024 · 6227562 · 6227562
1 parent 9e10936
commit 6227562
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh
@@ -126,6 +126,9 @@
   } else if (group_size == 2) {                              \
     constexpr size_t GROUP_SIZE = 2;                         \
     __VA_ARGS__                                              \
+  } else if (group_size == 3) {                              \
+    constexpr size_t GROUP_SIZE = 3;                         \
+    __VA_ARGS__                                              \
   } else if (group_size == 4) {                              \
     constexpr size_t GROUP_SIZE = 4;                         \
     __VA_ARGS__                                              \