From cb8f3d0ec65f51e879ccce48dab786ba239f150e Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Wed, 25 Mar 2026 16:00:01 -0700 Subject: [PATCH 1/2] First commit --- include/flashinfer/utils.cuh | 9 +++++++++ tests/attention/test_batch_decode_kernels.py | 8 ++++---- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index c7edf5ab57..3489079443 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -150,6 +150,15 @@ } else if (group_size == 8) { \ constexpr size_t GROUP_SIZE = 8; \ __VA_ARGS__ \ + } else if (group_size == 16) { \ + constexpr size_t GROUP_SIZE = 16; \ + __VA_ARGS__ \ + } else if (group_size == 32) { \ + constexpr size_t GROUP_SIZE = 32; \ + __VA_ARGS__ \ + } else if (group_size == 64) { \ + constexpr size_t GROUP_SIZE = 64; \ + __VA_ARGS__ \ } else { \ std::ostringstream err_msg; \ err_msg << "Unsupported group_size: " << group_size; \ diff --git a/tests/attention/test_batch_decode_kernels.py b/tests/attention/test_batch_decode_kernels.py index 39e736306a..5b0de666a2 100644 --- a/tests/attention/test_batch_decode_kernels.py +++ b/tests/attention/test_batch_decode_kernels.py @@ -62,7 +62,7 @@ def warmup_jit(): @pytest.mark.parametrize("batch_size", [12, 17, 128]) @pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384]) @pytest.mark.parametrize("page_size", [1, 8, 16]) -@pytest.mark.parametrize("num_kv_heads", [4]) +@pytest.mark.parametrize("num_kv_heads", [2, 4]) @pytest.mark.parametrize("num_qo_heads", [4, 32]) @pytest.mark.parametrize("head_dim", [128, 256]) @pytest.mark.parametrize("kv_layout", ["NHD"]) @@ -199,7 +199,7 @@ def test_batch_decode_with_paged_kv_cache( @pytest.mark.parametrize("batch_size", [12, 17, 128]) @pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384]) @pytest.mark.parametrize("page_size", [1, 8, 16]) -@pytest.mark.parametrize("num_kv_heads", [4]) +@pytest.mark.parametrize("num_kv_heads", [2, 4]) @pytest.mark.parametrize("num_qo_heads", [4, 32]) @pytest.mark.parametrize("head_dim", [128, 256]) @pytest.mark.parametrize("kv_layout", ["NHD"]) @@ -358,7 +358,7 @@ def test_batch_decode_with_paged_kv_cache_with_fast_plan( @pytest.mark.parametrize("batch_size", [12, 17, 128]) @pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384]) @pytest.mark.parametrize("page_size", [1, 8, 16]) -@pytest.mark.parametrize("num_kv_heads", [4]) +@pytest.mark.parametrize("num_kv_heads", [2, 4]) @pytest.mark.parametrize("num_qo_heads", [4, 32]) @pytest.mark.parametrize("head_dim", [128, 256]) @pytest.mark.parametrize("kv_layout", ["NHD"]) @@ -496,7 +496,7 @@ def test_batch_decode_with_tuple_paged_kv_cache( @pytest.mark.parametrize("batch_size", [12, 17, 128]) @pytest.mark.parametrize("kv_len", [54, 2048, 16384]) @pytest.mark.parametrize("page_size", [1, 8, 16]) -@pytest.mark.parametrize("num_kv_heads", [4]) +@pytest.mark.parametrize("num_kv_heads", [2, 4]) @pytest.mark.parametrize("num_qo_heads", [4, 32]) @pytest.mark.parametrize("head_dim", [128, 256]) @pytest.mark.parametrize("kv_layout", ["NHD"]) From afacf807920bb41083109b857f7c6fd3b0f4c91b Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Wed, 25 Mar 2026 16:15:29 -0700 Subject: [PATCH 2/2] Remove group size 64 --- include/flashinfer/utils.cuh | 3 --- 1 file changed, 3 deletions(-) diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index 3489079443..f911527829 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -156,9 +156,6 @@ } else if (group_size == 32) { \ constexpr size_t GROUP_SIZE = 32; \ __VA_ARGS__ \ - } else if (group_size == 64) { \ - constexpr size_t GROUP_SIZE = 64; \ - __VA_ARGS__ \ } else { \ std::ostringstream err_msg; \ err_msg << "Unsupported group_size: " << group_size; \