reduce the number of unit tests

Signed-off-by: Siyuan Fu <[email protected]>

Author: IwakuraRein

tests/attention/test_trtllm_gen_attention.py

@@ -339,40 +339,7 @@ def unpack_compare_nvfp4(
     return output_unpacked, output_ref
 
 
-@pytest.mark.parametrize("kv_layout", ["HND", "NHD"])
-@pytest.mark.parametrize(
-    "batch_size,page_size,num_kv_heads,head_grp_size",
-    [
-        (4, 16, 2, 1),
-        (4, 32, 4, 5),
-        (4, 64, 4, 8),
-        (128, 16, 2, 5),
-        (128, 32, 4, 1),
-        (128, 64, 2, 8),
-        (256, 16, 4, 8),
-        (256, 32, 2, 8),
-        (256, 64, 4, 1),
-        (256, 64, 4, 5),
-    ],
-)
-@pytest.mark.parametrize("window_left", [-1])  # todo(Siyuan): add 127 window_left
-@pytest.mark.parametrize(
-    "q_dtype,kv_dtype,o_dtype",
-    [
-        ("bf16", "bf16", "bf16"),
-        ("fp16", "fp16", "fp16"),
-        ("fp8", "fp8", "bf16"),
-        ("fp8", "fp8", "fp16"),
-        ("fp8", "fp8", "fp8"),
-        ("fp8", "fp8", "nvfp4"),
-    ],
-)
-@pytest.mark.parametrize("enable_pdl", [True, False, None])
-@pytest.mark.parametrize("enable_sink", [True, False])
-@pytest.mark.parametrize("max_q_len", [511])
-@pytest.mark.parametrize("max_kv_len", [2047])
-@pytest.mark.parametrize("device_scale", [True, False])
-def test_trtllm_batch_prefill(
+def _test_trtllm_batch_prefill(
     kv_layout,
     batch_size,
     page_size,
@@ -580,6 +547,71 @@ def test_trtllm_batch_prefill(
         assert (workspace_buffer[: 8192 * 256 * 4].cpu().numpy() == 0).all()
 
 
+@pytest.mark.parametrize("kv_layout", ["HND", "NHD"])
+@pytest.mark.parametrize(
+    "batch_size,page_size,num_kv_heads,head_grp_size",
+    [
+        (4, 16, 2, 1),
+        (4, 32, 4, 5),
+        (4, 64, 4, 8),
+        (128, 16, 2, 5),
+        (128, 32, 4, 1),
+        (128, 64, 2, 8),
+        (256, 16, 4, 8),
+        (256, 32, 2, 8),
+        (256, 64, 4, 1),
+        (256, 64, 4, 5),
+    ],
+)
+@pytest.mark.parametrize("window_left", [-1])  # todo(Siyuan): add 127 window_left
+@pytest.mark.parametrize(
+    "q_dtype,kv_dtype,o_dtype",
+    [
+        ("bf16", "bf16", "bf16"),
+        ("fp16", "fp16", "fp16"),
+        ("fp8", "fp8", "bf16"),
+        ("fp8", "fp8", "fp16"),
+        ("fp8", "fp8", "fp8"),
+        ("fp8", "fp8", "nvfp4"),
+    ],
+)
+@pytest.mark.parametrize("enable_pdl", [None])
+@pytest.mark.parametrize("enable_sink", [True, False])
+@pytest.mark.parametrize("max_q_len", [511])
+@pytest.mark.parametrize("max_kv_len", [2047])
+def test_trtllm_batch_prefill(
+    kv_layout,
+    batch_size,
+    page_size,
+    num_kv_heads,
+    head_grp_size,
+    window_left,
+    q_dtype,
+    o_dtype,
+    kv_dtype,
+    enable_pdl,
+    enable_sink,
+    max_q_len,
+    max_kv_len,
+):
+    _test_trtllm_batch_prefill(
+        kv_layout,
+        batch_size,
+        page_size,
+        num_kv_heads,
+        head_grp_size,
+        window_left,
+        q_dtype,
+        o_dtype,
+        kv_dtype,
+        enable_pdl,
+        enable_sink,
+        max_q_len,
+        max_kv_len,
+        kv_dtype == "fp8",
+    )
+
+
 @pytest.mark.parametrize("kv_layout", ["HND", "NHD"])
 @pytest.mark.parametrize(
     "batch_size,page_size,num_kv_heads,head_grp_size",
@@ -613,7 +645,7 @@ def test_trtllm_batch_prefill_bs1(
     max_q_len,
     max_kv_len,
 ):
-    test_trtllm_batch_prefill(
+    _test_trtllm_batch_prefill(
         kv_layout,
         batch_size,
         page_size,
@@ -966,7 +998,6 @@ def _test_trtllm_batch_decode(
 @pytest.mark.parametrize("enable_sink", [True, False])
 @pytest.mark.parametrize("max_in_kv_len", [110])
 @pytest.mark.parametrize("head_dim", [128])
-@pytest.mark.parametrize("device_scale", [True, False])
 def test_trtllm_batch_decode(
     backend,
     kv_layout,
@@ -983,7 +1014,6 @@ def test_trtllm_batch_decode(
     enable_sink,
     max_in_kv_len,
     head_dim,
-    device_scale,
 ):
     # General set of tests for trtllm-gen decode
     _test_trtllm_batch_decode(
@@ -1002,7 +1032,7 @@ def test_trtllm_batch_decode(
         enable_sink,
         max_in_kv_len,
         head_dim,
-        device_scale,
+        kv_dtype == "fp8",
     )
 
 
@@ -1024,6 +1054,7 @@ def test_trtllm_batch_decode(
 @pytest.mark.parametrize("enable_sink", [False])
 @pytest.mark.parametrize("max_in_kv_len", [8192])
 @pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("device_scale", [True, False])
 def test_trtllm_batch_decode_bs1(
     kv_layout,
     batch_size,
@@ -1039,6 +1070,7 @@ def test_trtllm_batch_decode_bs1(
     enable_sink,
     max_in_kv_len,
     head_dim,
+    device_scale,
 ):
     # Small number of test cases for batch size 1
     pytest.xfail("trtllm-gen decode gets incorrect output with bs1")
@@ -1058,7 +1090,7 @@ def test_trtllm_batch_decode_bs1(
         enable_sink,
         max_in_kv_len,
         head_dim,
-        False,
+        device_scale,
     )
 
 
@@ -1091,6 +1123,7 @@ def test_trtllm_batch_decode_bs1(
 @pytest.mark.parametrize("enable_sink", [False])
 @pytest.mark.parametrize("max_in_kv_len", [110])
 @pytest.mark.parametrize("head_dim", [256])
+@pytest.mark.parametrize("device_scale", [True, False])
 def test_trtllm_batch_decode_head_dim_256(
     kv_layout,
     batch_size,
@@ -1106,6 +1139,7 @@ def test_trtllm_batch_decode_head_dim_256(
     enable_sink,
     max_in_kv_len,
     head_dim,
+    device_scale,
 ):
     # Small number of test cases for head_dim = 256
     pytest.xfail("trtllm-gen decode gets incorrect output with head_dim = 256")
@@ -1125,7 +1159,7 @@ def test_trtllm_batch_decode_head_dim_256(
         enable_sink,
         max_in_kv_len,
         head_dim,
-        True,
+        device_scale,
     )
 
 
@@ -1151,6 +1185,7 @@ def test_trtllm_batch_decode_head_dim_256(
 @pytest.mark.parametrize("enable_sink", [False])
 @pytest.mark.parametrize("max_in_kv_len", [4096, 8192, 16384, 32768, 65536, 131072])
 @pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("device_scale", [True, False])
 def test_trtllm_batch_decode_long_sequence_length(
     kv_layout,
     batch_size,
@@ -1166,6 +1201,7 @@ def test_trtllm_batch_decode_long_sequence_length(
     enable_sink,
     max_in_kv_len,
     head_dim,
+    device_scale,
 ):
     # Small number of test cases for long sequence length
     _test_trtllm_batch_decode(
@@ -1184,7 +1220,7 @@ def test_trtllm_batch_decode_long_sequence_length(
         enable_sink,
         max_in_kv_len,
         head_dim,
-        False,
+        device_scale,
     )