From 702ced6535adf14a77a10c7f530dc98568cfd2e1 Mon Sep 17 00:00:00 2001 From: Lequn Chen Date: Fri, 25 Oct 2024 19:40:49 +0000 Subject: [PATCH] Add group size 3 to GQA dispatch --- include/flashinfer/utils.cuh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index 23f3c0f00..4efda679f 100644 --- 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__ \