Fix the assumption that pages_per_seq is already a multiple of num_kv_pages_per_blk.

PiperOrigin-RevId: 735851301

Author: bythew3i

jax/experimental/pallas/ops/tpu/ragged_paged_attention.py

@@ -43,14 +43,22 @@ def __init__(
   ):
     self._vmem_buf = vmem_buf
     seq_id, kv_pages_start = offset
-    self._async_copies = [
-        pltpu.make_async_copy(
-            pages_hbm_ref.at[page_indices_ref[seq_id, kv_pages_start + i]],
-            vmem_buf.at[i],
-            sem,
-        )
-        for i in range(vmem_buf.shape[0])
-    ]
+    pages_per_seq = page_indices_ref.shape[1]
+    self._async_copies = []
+    # TODO(jevinjiang): Only fetch dynamic shape in need! This will insert
+    # a bunch of if-ops. Check the performance when we have benchmarking setup.
+    for i in range(vmem_buf.shape[0]):
+      page_idx = kv_pages_start + i
+      page_idx = jax.lax.select(
+          page_idx < pages_per_seq, page_idx, pages_per_seq - 1
+      )
+      self._async_copies.append(
+          pltpu.make_async_copy(
+              pages_hbm_ref.at[page_indices_ref[seq_id, page_idx]],
+              vmem_buf.at[i],
+              sem,
+          )
+      )
 
   def start(self):
     """Starts the async copies."""

tests/pallas/tpu_ragged_paged_attention_test.py

@@ -64,10 +64,6 @@ def _test_ragged_paged_attention(
     max_num_seq = max(len(seq_lens), max_num_seq)
     max_kv_len = max(kv_lens)
     pages_per_seq = ceil_div(max_kv_len, page_size)
-    pages_per_seq = (
-        ceil_div(pages_per_seq, num_kv_pages_per_block)
-        * num_kv_pages_per_block
-    )
     num_q_heads, num_kv_heads = num_heads
 
     cu_q_lens = jnp.array(cu_q_lens, dtype=jnp.int32)
@@ -130,8 +126,8 @@ def _test_ragged_paged_attention(
         num_seqs=num_seqs,
     )
     tols = {
-        "float32": 1e-1,
-        "bfloat16": 2e-1,
+        "float32": 0.15,
+        "bfloat16": 0.2,
     }
     tol = tols[jnp.dtype(dtype).name]
     self.assertAllClose(output, expected, atol=tol, rtol=tol)