Add padding ragged paged attention test (#8741)

Author: vanbasten23

test/test_pallas.py

@@ -87,6 +87,9 @@ def _pagedattention_generate_qkv(
       q = torch.randn(batch_size, query_len, num_heads, head_dim, dtype=dtype)
     return q, k_pages, v_pages, page_indices
 
+  def _round_up_closest_multiple_of(self, x, base):
+    return (x + base - 1) // base * base
+
   def _ragged_pagedattention_generate_qkv(
       self,
       seq_lens,
@@ -95,6 +98,8 @@ def _ragged_pagedattention_generate_qkv(
       page_size,
       num_pages,
       dtype=torch.float32,
+      num_queries_per_block=None,
+      pad_num_q_tokens=False,
   ):
     num_seqs = len(seq_lens)
     # Make sure the q_len is no longer than the kv_len. For example,
@@ -106,7 +111,10 @@ def _ragged_pagedattention_generate_qkv(
       assert cur_q_len <= cur_kv_len, f"cur_q_len must be less than or equal to cur_kv_len. Got {cur_q_len} and {cur_kv_len}"
 
     query_lens = [seq_len[0] for seq_len in seq_lens]
-    num_q_tokens = sum(query_lens)
+    actual_num_q_tokens = sum(query_lens)
+    num_q_tokens = self._round_up_closest_multiple_of(
+        actual_num_q_tokens,
+        num_queries_per_block) if pad_num_q_tokens else actual_num_q_tokens
     kv_lens = torch.tensor([seq_len[1] for seq_len in seq_lens],
                            dtype=torch.int32)
     num_q_heads = num_heads[0]
@@ -727,34 +735,28 @@ def test_ragged_paged_attention_wrapper_without_dynamo(self):
         torch.allclose(
             output.cpu(), nonkernel_output.cpu(), atol=2e-1, rtol=1e-2))
 
-  @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
-                   "This test only works on TPUv4+.")
-  def test_ragged_paged_attention_wrapper_with_dynamo(self):
-    seq_lens = [
-        (1, 1328),
-        (5, 18),
-        (1, 129),
-        (120, 229),
-        (1, 122),  # first physical q block
-        (1, 64),
-        (32, 100),
-        (250, 463),
-        (1, 18),
-        (1, 17),
-        (99, 123)
-    ]  # last 3 physical q blocks [(q_len, kv_len),...]
-    num_heads = (4, 4)
-    head_dim = 128
-    dtype = torch.float32
-    page_size = 16
-    num_pages = 32768
+  def _verify_ragged_paged_attention_with_dynamo(
+      self,
+      seq_lens,
+      num_heads,
+      head_dim,
+      page_size,
+      num_pages,
+      dtype,
+      num_kv_pages_per_block,
+      num_queries_per_block,
+      pad_num_q_tokens=False,
+  ):
     num_seqs = len(seq_lens)
-    num_kv_pages_per_block = 128
-    num_queries_per_block = 8
-    block_kv_size = 256
-
     q, k_pages, v_pages, page_indices, cu_q_lens, kv_lens = self._ragged_pagedattention_generate_qkv(
-        seq_lens, num_heads, head_dim, page_size, num_pages, dtype=dtype)
+        seq_lens,
+        num_heads,
+        head_dim,
+        page_size,
+        num_pages,
+        dtype=dtype,
+        num_queries_per_block=num_queries_per_block,
+        pad_num_q_tokens=pad_num_q_tokens)
 
     q_xla = q.to("xla")
     k_pages_xla = k_pages.to("xla")
@@ -783,7 +785,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
     compiled_paged_attention = torch.compile(
         ragged_paged_attention_wrapper, backend="openxla")
 
-    output = compiled_paged_attention(
+    kernel_output = compiled_paged_attention(
         q_xla,
         k_pages_xla,
         v_pages_xla,
@@ -809,9 +811,117 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
         use_kernel=False,
     )
 
-    self.assertTrue(
-        torch.allclose(
-            output.cpu(), nonkernel_output.cpu(), atol=2e-1, rtol=1e-2))
+    kernel_output_cpu = kernel_output.cpu()
+    nonkernel_output_cpu = nonkernel_output.cpu()
+    self.assertEqual(kernel_output_cpu.shape, nonkernel_output_cpu.shape)
+    self.assertEqual(kernel_output_cpu.dtype, nonkernel_output_cpu.dtype)
+
+    q_jax = jnp.array(q.numpy(), dtype=jnp.float32)
+    k_pages_jax = jnp.array(k_pages.numpy(), dtype=jnp.float32)
+    v_pages_jax = jnp.array(v_pages.numpy(), dtype=jnp.float32)
+    kv_lens_jax = jnp.array(kv_lens.numpy(), dtype=jnp.int32)
+    page_indices_jax = jnp.array(page_indices.numpy(), dtype=jnp.int32)
+    cu_q_lens_jax = jnp.array(cu_q_lens.numpy(), dtype=jnp.int32)
+
+    from torch_xla.experimental.pallas_kernels.ragged_paged_attention_kernel import ragged_paged_attention as jax_ragged_paged_attention
+    jax_kernel_output = torch.from_numpy(
+        np.array(
+            jax_ragged_paged_attention(
+                q_jax,
+                k_pages_jax,
+                v_pages_jax,
+                kv_lens_jax,
+                page_indices_jax,
+                cu_q_lens_jax,
+                num_seqs=num_seqs,
+                num_kv_pages_per_block=num_kv_pages_per_block,
+                num_queries_per_block=num_queries_per_block,
+            )[1]))
+    jax_kernel_output_cpu = jax_kernel_output.cpu()
+
+    if pad_num_q_tokens:
+      actual_num_q_tokens = cu_q_lens[num_seqs]
+      self.assertTrue(
+          torch.allclose(
+              kernel_output_cpu[:actual_num_q_tokens],
+              nonkernel_output_cpu[:actual_num_q_tokens],
+              atol=2e-1,
+              rtol=1e-2))
+      self.assertTrue(
+          torch.allclose(
+              kernel_output_cpu[:actual_num_q_tokens],
+              jax_kernel_output_cpu[:actual_num_q_tokens],
+              atol=2e-1,
+              rtol=1e-2))
+    else:
+      self.assertTrue(
+          torch.allclose(
+              kernel_output_cpu, nonkernel_output_cpu, atol=2e-1, rtol=1e-2))
+      self.assertTrue(
+          torch.allclose(
+              kernel_output_cpu, jax_kernel_output_cpu, atol=2e-1, rtol=1e-2))
+
+  @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
+                   "This test only works on TPUv4+.")
+  def test_ragged_paged_attention_wrapper_no_query_padding_with_dynamo(self):
+    seq_lens = [
+        (1, 1328),
+        (5, 18),
+        (1, 129),
+        (120, 229),
+        (1, 122),  # first physical q block
+        (1, 64),
+        (32, 100),
+        (250, 463),
+        (1, 18),
+        (1, 17),
+        (99, 123)
+    ]  # last 3 physical q blocks [(q_len, kv_len),...]
+    num_heads = (4, 4)
+    head_dim = 128
+    dtype = torch.float32
+    page_size = 16
+    num_pages = 32768
+
+    self._verify_ragged_paged_attention_with_dynamo(
+        seq_lens,
+        num_heads,
+        head_dim,
+        page_size,
+        num_pages,
+        dtype,
+        num_kv_pages_per_block=128,
+        num_queries_per_block=8,
+    )
+
+  @parameterized.product(
+      seq_lens=[[(1, 1328), (5, 18), (500, 563)]],
+      num_queries_per_block=[16, 64, 128],
+  )
+  @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
+                   "This test only works on TPUv4+.")
+  def test_ragged_paged_attention_wrapper_with_query_padding_with_dynamo(
+      self,
+      seq_lens,
+      num_queries_per_block,
+  ):
+    num_heads = (4, 4)
+    head_dim = 128
+    dtype = torch.float32
+    page_size = 16
+    num_pages = 32768
+
+    self._verify_ragged_paged_attention_with_dynamo(
+        seq_lens,
+        num_heads,
+        head_dim,
+        page_size,
+        num_pages,
+        dtype,
+        num_kv_pages_per_block=128,
+        num_queries_per_block=num_queries_per_block,
+        pad_num_q_tokens=True,
+    )
 
   @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
                    "This test only works on TPUv4+.")

test/test_ragged_paged_attention_kernel.py

@@ -1,16 +1,16 @@
 from typing import List, Optional, Tuple
+import sys
+import unittest
 
-from absl.testing import absltest
 from absl.testing import parameterized
+from absl.testing import absltest
 import jax
 from jax._src import test_util as jtu
 from jax.experimental.pallas.ops.tpu.paged_attention import quantization_utils
 from torch_xla.experimental.pallas_kernels.ragged_paged_attention_kernel import ragged_paged_attention, make_sequence_metadata, DEFAULT_MASK_VALUE
 import jax.numpy as jnp
 import numpy as np
 
-jax.config.parse_flags_with_absl()
-
 ATOL_FP32 = 2e-1
 
 
@@ -29,6 +29,7 @@ def _ref_ragged_paged_attention(
   assert num_q_heads % num_kv_heads == 0, "num_q_heads % num_kv_heads !=0."
   num_query_per_kv = num_q_heads // num_kv_heads
   start_idx = 0
+
   outputs: List[jax.Array] = []
   for i in range(num_seqs):
     cur_q_len = cu_q_lens[i + 1] - cu_q_lens[i]
@@ -72,11 +73,17 @@ def _ref_ragged_paged_attention(
     outputs.append(out)
     start_idx += cur_q_len
 
+  maybe_padded_num_q_tokens = queries.shape[0]
+  actual_num_tokens = cu_q_lens[num_seqs]
+  if actual_num_tokens < maybe_padded_num_q_tokens:
+    num_tokens_diff = maybe_padded_num_q_tokens - actual_num_tokens
+    outputs.append(
+        jnp.zeros(
+            (num_tokens_diff, num_q_heads, head_dim)).astype(outputs[0].dtype))
   return jnp.concatenate(outputs, axis=0)
 
 
-@jtu.with_config(jax_numpy_dtype_promotion="standard")
-class RaggedPagedAttentionKernelTest(jtu.JaxTestCase):
+class RaggedPagedAttentionKernelTest(parameterized.TestCase):
 
   def _verify_ragged_paged_attention(
       self,
@@ -88,6 +95,7 @@ def _verify_ragged_paged_attention(
       num_pages,
       num_kv_pages_per_block=128,
       num_queries_per_block=128,
+      pad_num_q_tokens=False,
   ):
     num_seqs = len(seq_lens)
     # Make sure the q_len is no longer than the kv_len. For example,
@@ -99,7 +107,11 @@ def _verify_ragged_paged_attention(
       assert cur_q_len <= cur_kv_len, f"cur_q_len must be less than or equal to cur_kv_len. Got {cur_q_len} and {cur_kv_len}"
 
     query_lens = [seq_len[0] for seq_len in seq_lens]
-    num_q_tokens = sum(query_lens)
+    actual_num_q_tokens = sum(query_lens)
+    # Caller(eg vLLM) may decide to pad the num_q_tokens.
+    num_q_tokens = self._round_up_closest_multiple_of(
+        actual_num_q_tokens,
+        num_queries_per_block) if pad_num_q_tokens else actual_num_q_tokens
     kv_lens = jnp.array([seq_len[1] for seq_len in seq_lens])
     num_q_heads = num_heads[0]
     num_kv_heads = num_heads[1]
@@ -115,6 +127,8 @@ def _verify_ragged_paged_attention(
         k3, (num_kv_heads, num_pages, page_size, head_dim), dtype=dtype)
 
     # Create a kv_lens: i32[num_tokens]
+    # Only the first num_seqs of kv_lens_with_paddings are meaningful
+    # [num_seqs:num_q_tokens] are padded value and are meaningless.
     kv_lens_with_paddings = [0] * num_q_tokens
     for i in range(num_seqs):
       kv_lens_with_paddings[i] = kv_lens[i]
@@ -182,8 +196,16 @@ def _verify_ragged_paged_attention(
       rtol = 1e-1
     else:
       self.fail(f'Unsupported dtype: {dtype}')
-    self.assertTrue(
-        jnp.allclose(actual_output, expected_output, atol=atol, rtol=rtol))
+    if pad_num_q_tokens:
+      self.assertTrue(
+          jnp.allclose(
+              actual_output[:actual_num_q_tokens],
+              expected_output[:actual_num_q_tokens],
+              atol=atol,
+              rtol=rtol))
+    else:
+      self.assertTrue(
+          jnp.allclose(actual_output, expected_output, atol=atol, rtol=rtol))
 
   def _round_up_closest_multiple_of(self, x, base):
     return (x + base - 1) // base * base
@@ -215,11 +237,12 @@ def test_paged_attention_basic(self,):
 
   @parameterized.product(
       seq_lens=[[(1, 1328), (5, 18), (506, 563)]],
-      num_heads=[(4, 4), (8, 2), (16, 2)],
+      num_heads=[(4, 4), (4, 2)],
       head_dim=[128, 256],
       dtype=(jnp.float32, jnp.bfloat16),
       page_size=[16, 32],
       num_pages=[32768, 2048],
+      num_queries_per_block=[16, 64, 128],
   )
   def test_paged_attention_varlen_comprehensive(
       self,
@@ -229,6 +252,7 @@ def test_paged_attention_varlen_comprehensive(
       dtype,
       page_size: int,
       num_pages: int,
+      num_queries_per_block: int,
   ):
     if jtu.is_device_tpu(version=4) and head_dim == 256 and page_size == 32:
       self.skipTest(
@@ -240,7 +264,42 @@ def test_paged_attention_varlen_comprehensive(
         page_size,
         dtype,
         num_pages,
-        num_queries_per_block=64,
+        num_queries_per_block=num_queries_per_block,
+        num_kv_pages_per_block=128,
+    )
+
+  @parameterized.product(
+      num_heads=[(4, 4), (4, 2)],
+      head_dim=[128, 256],
+      dtype=(jnp.float32, jnp.bfloat16),
+      page_size=[16, 32],
+      num_pages=[32768, 2048],
+      num_queries_per_block=[16, 64, 128],
+  )
+  def test_paged_attention_varlen_with_padding_comprehensive(
+      self,
+      num_heads: Tuple[int, int],
+      head_dim: int,
+      dtype,
+      page_size: int,
+      num_pages: int,
+      num_queries_per_block: int,
+  ):
+    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.")
+    # If num_queries_per_block is 128, then num_tokens will be pad 6 to be the smallest multiple of 128.
+    seq_lens = [(1, 1328), (5, 18), (500, 563)]
+    self._verify_ragged_paged_attention(
+        seq_lens,
+        num_heads,
+        head_dim,
+        page_size,
+        dtype,
+        num_pages,
+        num_queries_per_block=num_queries_per_block,
+        num_kv_pages_per_block=128,
+        pad_num_q_tokens=True,
     )
 
   def test_paged_attention_mix_prefill_and_decode1(self,):
@@ -442,4 +501,5 @@ def test_make_sequence_metadata(self,):
 
 
 if __name__ == "__main__":
-  absltest.main(testLoader=jtu.JaxTestLoader())
+  test = unittest.main()
+  sys.exit(0 if test.result.wasSuccessful() else 1)