Fix ragged_paged_attention op signature (#8943)

Author: yaochengji

test/test_pallas.py

@@ -678,6 +678,7 @@ def ragged_paged_attention_wrapper(
           sliding_window=sliding_window,
           soft_cap=soft_cap,
           use_kernel=True,
+          max_model_len=2048,
           num_kv_pages_per_block=num_kv_pages_per_block,
           num_queries_per_block=num_queries_per_block,
       ):
@@ -692,6 +693,7 @@ def ragged_paged_attention_wrapper(
             sliding_window=sliding_window,
             soft_cap=soft_cap,
             use_kernel=use_kernel,
+            max_model_len=max_model_len,
             num_kv_pages_per_block=num_kv_pages_per_block,
             num_queries_per_block=num_queries_per_block,
         )
@@ -712,6 +714,7 @@ def ragged_paged_attention_wrapper(
         sliding_window=sliding_window,
         soft_cap=soft_cap,
         use_kernel=True,
+        max_model_len=2048,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
     )[:cu_q_lens[num_seqs]]
@@ -752,12 +755,12 @@ def ragged_paged_attention_wrapper(
 
     from torch_xla.experimental.pallas_kernels.ragged_paged_attention_v2 import ragged_paged_attention as jax_ragged_paged_attention
     from torch_xla.experimental.tuned_block_sizes import get_ragged_attention_tuned_block_size
+    max_model_len = 2048
     if num_kv_pages_per_block is None:
       assert num_queries_per_block is None
       token_num = q.shape[0]
       token_num, q_head_num, _ = q.shape
       kv_head_num = kv_pages[2] // 2
-      max_model_len = 2048
       num_kv_pages_per_block, num_queries_per_block = get_ragged_attention_tuned_block_size(
           q_head_num, kv_head_num, token_num, max_model_len)
     jax_kernel_output = torch.from_numpy(

torch_xla/experimental/custom_kernel.py

@@ -1681,7 +1681,7 @@ def non_xla_ragged_paged_attention(q, kv, attention_type):
 XLA_LIB.define(
     "ragged_paged_attention(Tensor q, Tensor kv_pages, Tensor kv_lens, Tensor page_indices, "
     "Tensor cu_q_lens, Tensor num_seqs, float sm_scale=1, int? sliding_window=None, "
-    "float? soft_cap=None, float? mask_value=None, bool use_kernel=True, "
+    "float? soft_cap=None, float? mask_value=None, bool use_kernel=True, int max_model_len=2048,"
     "int? num_kv_pages_per_block=None, int? num_queries_per_block=None, int? vmem_limit_bytes=None) -> Tensor",
 )
 
@@ -1699,6 +1699,7 @@ def ragged_paged_attention_xla(
     soft_cap: float | None = None,
     mask_value=None,
     use_kernel=True,
+    max_model_len=2048,
     # kernel tuning parameters
     num_kv_pages_per_block=None,
     num_queries_per_block=None,
@@ -1716,6 +1717,7 @@ def ragged_paged_attention_xla(
       soft_cap=soft_cap,
       mask_value=mask_value,
       use_kernel=use_kernel,
+      max_model_len=max_model_len,
       num_kv_pages_per_block=num_kv_pages_per_block,
       num_queries_per_block=num_queries_per_block,
       vmem_limit_bytes=vmem_limit_bytes)
@@ -1734,6 +1736,7 @@ def ragged_paged_attention_non_xla(
     soft_cap: float | None = None,
     mask_value=None,
     use_kernel=True,
+    max_model_len=2048,
     # kernel tuning parameters
     num_kv_pages_per_block=None,
     num_queries_per_block=None,