Remove dynamic grid (#8896)

Author: bythew3i

test/test_pallas.py

@@ -669,10 +669,6 @@ def _test_ragged_paged_attention(
     page_indices_xla = page_indices.to("xla")
     cu_q_lens_xla = cu_q_lens.to("xla")
     num_seqs_xla = torch.tensor([num_seqs], dtype=torch.int32).to("xla")
-    sliding_window = sliding_window
-    soft_cap = soft_cap
-    # Test mask_value
-    mask_value = None
 
     if use_dynamo:
 
@@ -686,7 +682,6 @@ def ragged_paged_attention_wrapper(
           sm_scale=sm_scale,
           sliding_window=sliding_window,
           soft_cap=soft_cap,
-          mask_value=mask_value,
           use_kernel=True,
           num_kv_pages_per_block=num_kv_pages_per_block,
           num_queries_per_block=num_queries_per_block,
@@ -701,7 +696,6 @@ def ragged_paged_attention_wrapper(
             sm_scale=sm_scale,
             sliding_window=sliding_window,
             soft_cap=soft_cap,
-            mask_value=mask_value,
             use_kernel=use_kernel,
             num_kv_pages_per_block=num_kv_pages_per_block,
             num_queries_per_block=num_queries_per_block,
@@ -722,7 +716,6 @@ def ragged_paged_attention_wrapper(
         sm_scale=sm_scale,
         sliding_window=sliding_window,
         soft_cap=soft_cap,
-        mask_value=mask_value,
         use_kernel=True,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
@@ -738,7 +731,6 @@ def ragged_paged_attention_wrapper(
         sm_scale=sm_scale,
         sliding_window=sliding_window,
         soft_cap=soft_cap,
-        mask_value=mask_value,
         use_kernel=False,
     )
 
@@ -778,7 +770,6 @@ def ragged_paged_attention_wrapper(
                 sm_scale=sm_scale,
                 sliding_window=sliding_window,
                 soft_cap=soft_cap,
-                mask_value=mask_value,
             )[:cu_q_lens[num_seqs]].astype(jnp.float32))).to(dtype)
     jax_kernel_output_cpu = jax_kernel_output.cpu()
 

torch_xla/experimental/custom_kernel.py

@@ -1015,14 +1015,8 @@ def ragged_paged_attention(
   )
 
   seq_buf_idx = torch.tensor([0, 0], dtype=torch.int32).to("xla")
-  num_q_blks = torch.tensor(
-      [(cu_q_lens[num_seqs[0]] + num_queries_per_block - 1) //
-       num_queries_per_block],
-      dtype=torch.int32).to("xla")
-
   output = torch_xla._XLAC._xla_tpu_custom_call(
       [
-          num_q_blks,  # dynamic grid
           kv_lens,
           page_indices,
           cu_q_lens,

torch_xla/experimental/pallas_kernels/ragged_paged_attention_v2.py

@@ -295,7 +295,9 @@ def prefetch_first_kv_blk():
 
   def is_cur_q_blk_needed(q_states):
     done, cur_seq_idx, _ = q_states
-    return jnp.logical_and(done == 0, cur_seq_idx < num_seqs)
+    should_run = jnp.logical_and(q_len_start < cu_q_lens_ref[num_seqs],
+                                 cur_seq_idx < num_seqs)
+    return jnp.logical_and(done == 0, should_run)
 
   def compute_with_cur_q_blk(q_states):
     done, cur_seq_idx, cur_buf_idx = q_states
@@ -640,14 +642,14 @@ def ragged_paged_attention(
   check_inputs_shapes(q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs)
   if mask_value is None:
     mask_value = DEFAULT_MASK_VALUE
-  _, num_q_heads, head_dim = q.shape
+  num_q, num_q_heads, head_dim = q.shape
   _, page_size, num_combined_kv_heads, _ = kv_pages.shape
   assert num_combined_kv_heads % 2 == 0
   num_kv_heads = num_combined_kv_heads // 2
   num_q_per_blk = num_queries_per_block
   num_kv_pages_per_blk = num_kv_pages_per_block
   num_q_heads_per_kv_head = num_q_heads // num_kv_heads
-  num_q_blks = cdiv(cu_q_lens[num_seqs[0]], num_q_per_blk)
+  num_q_blks = cdiv(num_q, num_q_per_blk)
   num_q_heads_per_blk, num_combined_kv_heads_per_blk = get_min_heads_per_blk(
       num_q_heads, num_combined_kv_heads, q.dtype, kv_pages.dtype)
   assert num_combined_kv_heads_per_blk % 2 == 0