[Kernel] Improve 2D Triton Attention Kernel

tests/kernels/attention/test_triton_unified_attention.py

@@ -22,6 +22,10 @@
 # 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]
+
 
 def ref_paged_attn(
     query: torch.Tensor,
@@ -92,6 +96,7 @@ 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)
 @torch.inference_mode()
 def test_triton_unified_attn(
     seq_lens: list[tuple[int, int]],
@@ -103,6 +108,7 @@ def test_triton_unified_attn(
     soft_cap: float | None,
     num_blocks: int,
     q_dtype: torch.dtype | None,
+    seq_threshold_3D: int,
 ) -> None:
     torch.set_default_device("cuda")
 
@@ -152,6 +158,28 @@ def test_triton_unified_attn(
         k_descale = torch.rand(scale_shape, dtype=torch.float32)
         v_descale = torch.rand(scale_shape, dtype=torch.float32)
 
+    num_queries_per_kv = num_query_heads // num_kv_heads
+    BLOCK_M = (
+        16 if num_queries_per_kv <= 16 else 1 << (num_queries_per_kv - 1).bit_length()
+    )  # next power of 2 value
+    BLOCK_Q = BLOCK_M // num_queries_per_kv
+
+    block_q_seq_boundaries_tensor = torch.empty(num_seqs + 1, dtype=torch.int32)
+    if max_query_len > 1:
+        block_q_seq_boundaries_tensor[0] = 0
+        block_q_seq_boundaries_tensor[1 : cu_query_lens.numel()].copy_(
+            cu_query_lens[1:]
+        )
+        block_q_seq_boundaries_tensor[1 : cu_query_lens.numel()].sub_(
+            cu_query_lens[:-1]
+        )
+        block_q_seq_boundaries_tensor[1 : cu_query_lens.numel()].add_(BLOCK_Q - 1)
+        block_q_seq_boundaries_tensor[1 : cu_query_lens.numel()].floor_divide_(BLOCK_Q)
+        block_q_seq_boundaries_tensor[: cu_query_lens.numel()].cumsum_(dim=0)
+        num_q_blocks = block_q_seq_boundaries_tensor[cu_query_lens.numel() - 1]
+    else:
+        num_q_blocks = len(seq_lens)
+
     unified_attention(
         q=maybe_quantized_query,
         k=maybe_quantized_key_cache,
@@ -169,6 +197,11 @@ def test_triton_unified_attn(
         q_descale=q_descale,
         k_descale=k_descale,
         v_descale=v_descale,
+        BLOCK_M=BLOCK_M,
+        BLOCK_Q=BLOCK_Q,
+        num_q_blocks=num_q_blocks,
+        block_q_seq_boundaries_tensor=block_q_seq_boundaries_tensor,
+        seq_threshold_3D=seq_threshold_3D,
     )
 
     ref_output = ref_paged_attn(