Set scoped vmem for paged attention (#8988)

Author: zpcore

test/test_ragged_paged_attention_kernel.py

@@ -254,6 +254,9 @@ def test_paged_attention_varlen_comprehensive(
       num_pages: int,
       num_queries_per_block: int,
   ):
+    if jtu.is_device_tpu(version=5, variant="e"):
+      self.skipTest(
+          "TPU v5e has small VMEM. It will run into VMEM OOM. Skip the test.")
     if jtu.is_device_tpu(version=4) and head_dim == 256 and page_size == 32:
       self.skipTest(
           "TPU v4 has small VMEM. It will run into VMEM OOM. Skip the test.")
@@ -285,6 +288,9 @@ def test_paged_attention_varlen_with_padding_comprehensive(
       num_pages: int,
       num_queries_per_block: int,
   ):
+    if jtu.is_device_tpu(version=5, variant="e"):
+      self.skipTest(
+          "TPU v5e has small VMEM. It will run into VMEM OOM. Skip the test.")
     if jtu.is_device_tpu(version=4) and head_dim == 256 and page_size == 32:
       self.skipTest(
           "TPU v4 has small VMEM. It will run into VMEM OOM. Skip the test.")

torch_xla/experimental/pallas_kernels/ragged_paged_attention_kernel.py

@@ -961,7 +961,9 @@ def next_kv_blk_page_indices_index_map(kv_head_idx, logical_q_blk_idx,
               "arbitrary",
               "arbitrary",
               "arbitrary",
-          )),
+          ),
+          vmem_limit_bytes=64 * 1024 * 1024,
+      ),
       out_shape=out_shape,
   )
   buffer_index = jnp.zeros((1,), jnp.int32)