Change num_seqs type from int to torch.Tensor (#8736)

Author: vanbasten23

test/test_pallas.py

@@ -682,6 +682,7 @@ def test_ragged_paged_attention_wrapper_without_dynamo(self):
     kv_lens_xla = kv_lens.to("xla")
     page_indices_xla = page_indices.to("xla")
     cu_q_lens_xla = cu_q_lens.to("xla")
+    num_seqs_xla = torch.tensor(num_seqs).to('xla')
 
     output = ragged_paged_attention(
         q_xla,
@@ -690,7 +691,7 @@ def test_ragged_paged_attention_wrapper_without_dynamo(self):
         kv_lens_xla,
         page_indices_xla,
         cu_q_lens_xla,
-        num_seqs=num_seqs,
+        num_seqs=num_seqs_xla,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=True)
@@ -702,7 +703,7 @@ def test_ragged_paged_attention_wrapper_without_dynamo(self):
         kv_lens_xla,
         page_indices_xla,
         cu_q_lens_xla,
-        num_seqs=num_seqs,
+        num_seqs=num_seqs_xla,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=False)
@@ -764,6 +765,7 @@ def _verify_ragged_paged_attention_with_dynamo(
     kv_lens_xla = kv_lens.to("xla")
     page_indices_xla = page_indices.to("xla")
     cu_q_lens_xla = cu_q_lens.to("xla")
+    num_seqs_xla = torch.tensor(num_seqs).to("xla")
 
     def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
                                        page_indices, cu_q_lens, num_seqs,
@@ -792,7 +794,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
         kv_lens_xla,
         page_indices_xla,
         cu_q_lens_xla,
-        num_seqs=num_seqs,
+        num_seqs=num_seqs_xla,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=True,
@@ -805,7 +807,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
         kv_lens_xla,
         page_indices_xla,
         cu_q_lens_xla,
-        num_seqs=num_seqs,
+        num_seqs=num_seqs_xla,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=False,

torch_xla/experimental/custom_kernel.py

@@ -780,13 +780,14 @@ def ragged_paged_attention(
     kv_lens,  # i32[num_tokens]
     page_indices,  # i32[num_tokens, pages_per_sequence]
     cu_q_lens,  # i32[num_tokens + 1]
-    num_seqs,  # int
+    num_seqs,  # i32[]
     num_kv_pages_per_block,
     num_queries_per_block,
     use_kernel=True,
     # TODO(jevinjiang, xiowei): add attn_logits_soft_cap.
     # attn_logits_soft_cap: float | None = None,
 ):  # [batch_size, query_len, num_heads, head_dim]:
+  num_seqs = num_seqs.item()
   assert len(q.shape) == 3, "q should have 3 dimensions."
   if not use_kernel:
     return _ragged_paged_attention_nonkernel(
@@ -1541,15 +1542,15 @@ def multi_queries_paged_attention_non_xla(q: torch.Tensor,
 
 
 XLA_LIB.define(
-    "ragged_paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor kv_lens, Tensor page_indices, Tensor cu_q_lens, int num_seqs, int num_kv_pages_per_block, int num_queries_per_block, bool use_kernel) -> Tensor",
+    "ragged_paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor kv_lens, Tensor page_indices, Tensor cu_q_lens, Tensor num_seqs, int num_kv_pages_per_block, int num_queries_per_block, bool use_kernel) -> Tensor",
 )
 
 
 @impl(XLA_LIB, "ragged_paged_attention", "XLA")
 def ragged_paged_attention_xla(q: torch.Tensor, k_pages: torch.Tensor,
                                v_pages: torch.Tensor, kv_lens: torch.Tensor,
                                page_indices: torch.Tensor,
-                               cu_q_lens: torch.Tensor, num_seqs: int,
+                               cu_q_lens: torch.Tensor, num_seqs: torch.Tensor,
                                num_kv_pages_per_block: int,
                                num_queries_per_block: int, use_kernel: bool):
   return ragged_paged_attention(q, k_pages, v_pages, kv_lens, page_indices,
@@ -1561,8 +1562,8 @@ def ragged_paged_attention_xla(q: torch.Tensor, k_pages: torch.Tensor,
 def ragged_paged_attention_non_xla(
     q: torch.Tensor, k_pages: torch.Tensor, v_pages: torch.Tensor,
     kv_lens: torch.Tensor, page_indices: torch.Tensor, cu_q_lens: torch.Tensor,
-    num_seqs: int, num_kv_pages_per_block: int, num_queries_per_block: int,
-    use_kernel: bool):
+    num_seqs: torch.Tensor, num_kv_pages_per_block: int,
+    num_queries_per_block: int, use_kernel: bool):
   return non_xla_attetion(q, k_pages, v_pages, "paged")