add back missing ut

Author: yzh119

tests/attention/test_batch_prefill_kernels.py

@@ -51,8 +51,8 @@ def warmup_jit():
 @pytest.mark.parametrize("kv_len", [54, 97, 512, 2048])
 @pytest.mark.parametrize("qo_len", [37, 17, 127, 577])
 @pytest.mark.parametrize("page_size", [1, 5, 16])
-@pytest.mark.parametrize("num_kv_heads", [32])
-@pytest.mark.parametrize("num_qo_heads", [32])
+@pytest.mark.parametrize("num_kv_heads", [4])
+@pytest.mark.parametrize("num_qo_heads", [4, 32])
 @pytest.mark.parametrize("head_dim", [128, 256])
 @pytest.mark.parametrize("causal", [False, True])
 @pytest.mark.parametrize("kv_layout", ["NHD"])