[Kernel] Separate Triton Attention Kernel Launches for Prefill and Decode for FULL CUDA Graph mode

tests/kernels/attention/test_triton_unified_attention.py

@@ -22,6 +22,12 @@
 # one value small enough to test the schema op check
 NUM_BLOCKS = [32768, 2048]
 
+# 0: use 2D kernel for decode
+# 8: use 3D kernel for decode
+SEQ_THRESHOLD_3D_VALUES = [0, 8]
+
+SPLIT_LAUNCH_VALUES = [False, True]
+
 
 def ref_paged_attn(
     query: torch.Tensor,
@@ -82,7 +88,12 @@ def ref_paged_attn(
 
 
 @pytest.mark.parametrize(
-    "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
+    "seq_lens",
+    [
+        [(1, 1328), (5, 18), (129, 463)],  # mixed batch
+        [(1, 523), (1, 37), (1, 2011)],  # decode-only batch
+        [(5, 18), (129, 463)],  # prefill-only batch
+    ],
 )
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
 @pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -92,6 +103,8 @@ def ref_paged_attn(
 @pytest.mark.parametrize("soft_cap", [None, 50.0])
 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
 @pytest.mark.parametrize("q_dtype", QDTYPES)
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
+@pytest.mark.parametrize("split_launch", SPLIT_LAUNCH_VALUES)
 @torch.inference_mode()
 def test_triton_unified_attn(
     seq_lens: list[tuple[int, int]],
@@ -103,6 +116,8 @@ def test_triton_unified_attn(
     soft_cap: float | None,
     num_blocks: int,
     q_dtype: torch.dtype | None,
+    seq_threshold_3D: int,
+    split_launch: bool,
 ) -> None:
     torch.set_default_device("cuda")
 
@@ -152,6 +167,8 @@ def test_triton_unified_attn(
         k_descale = torch.rand(scale_shape, dtype=torch.float32)
         v_descale = torch.rand(scale_shape, dtype=torch.float32)
 
+    num_decodes = num_seqs if max_query_len == 1 else query_lens.count(1)
+
     unified_attention(
         q=maybe_quantized_query,
         k=maybe_quantized_key_cache,
@@ -161,6 +178,7 @@ def test_triton_unified_attn(
         seqused_k=kv_lens,
         max_seqlen_q=max_query_len,
         max_seqlen_k=max_kv_len,
+        num_decodes=num_decodes,
         softmax_scale=scale,
         causal=True,
         window_size=window_size,
@@ -169,6 +187,8 @@ def test_triton_unified_attn(
         q_descale=q_descale,
         k_descale=k_descale,
         v_descale=v_descale,
+        seq_threshold_3D=seq_threshold_3D,
+        split_launch=split_launch,
     )
 
     ref_output = ref_paged_attn(